diff --git a/CHANGELOG.md b/CHANGELOG.md index 3bf19d05..54cc6f84 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,22 +2,28 @@ # CUTLASS 4.x -## [4.3.0](https://github.com/NVIDIA/cutlass/tree/main) (2025-10-20) +## [4.3.0](https://github.com/NVIDIA/cutlass/releases/tag/v4.3.0) (2025-11-21) ### CuTe DSL +* New features: + - Supported Apache [TVM-FFI](https://tvm.apache.org/ffi/index.html) for further reduced host runtime overhead for JIT functions, better PyTorch and ML frameworks interopability + - Added fake tensor and stream to decouple compile jit function with "from_dlpack" flow. Now we no longer require users to have real tensor when compile jit function. + - Added FastDivmodDivisor with Python operator overloads, new APIs, Cute dialect integration, and optimized static tile scheduler performance for faster index mapping. + - Added l2 cache evict priority for tma related ops. Users could do fine-grain l2 cache control. * Debuggability improvements: - - Supported source location tracking for DSL APIs - - Supported dumping PTX and CUBIN code + - Supported source location tracking for DSL APIs (Allow tools like ``nsight`` profiling to correlate perf metrics with Python source code) + - Supported dumping PTX and CUBIN code: [Hello World Example](https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/notebooks/hello_world.ipynb) * More examples and notebooks to get started with CuTe DSL: - - Improved performance of elementwise kernel (https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/ampere/elementwise_apply.py): + - Improved performance of [elementwise example](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/ampere/elementwise_apply.py): + Generalize code to handle list of input tensors + Generalize TV layout computation to handle different data types - - Demonstrate the new Pipeline APIs in [Blackwell SM100 persistent dense GEMM with static scheduling](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py): - + New Pipeline API `PipelineProducer` and `PipelineConsumer` to simplify code (no more explicit pipeline state management) - - Separate epilogue code for non-TMA and TMA implementation - + Note that the updates simplifies the codes but existing APIs still work and are supported - - [Basic Blackwell SM100 GEMM with decent performance](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py) - + Simple tutorial achieves 84% SOL performance with MNK 8K + - Improved [Blackwell SM100 persistent dense GEMM with static scheduling](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py): + + To demonstrate usage of new Pipeline APIs `PipelineProducer` and `PipelineConsumer` to simplify code without explicit pipeline state management (Exiting APIs are still maintained) + + Separated epilogue code for non-TMA and TMA implementation + - [Tutorial for Blackwell GEMM: Basic Blackwell SM100 GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/tutorial_gemm) + + [Baseline Blackwell GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py) achieves 84% SOL performance with MNK 8K + + More examples are coming for demo of optimization: `Baseline + X` + - [Tutorial for Async Pipeline API](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/async_pipeline.ipynb) - Reworked [elementwise add notebook](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/elementwise_add.ipynb) with more details and detailed explanation about TV layout + Updated implementation to handle general data type and multiple inputs + Updated explanation for TV layout in simpler language @@ -44,26 +50,49 @@ - ``PipelineAsync`` - ``SmemAllocator`` - Fixed documentation for ``pipeline``, ``utils`` and ``cute.math`` + - Added overlapping accumulator optimization for block tile N = 256 case for better epilogue latency hiding 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). + - Fixed TensorSSA.__getitem__ indexing to match CuTe's indexing convention + - Fixed an issue with cutlass.max and cutlass.min + - Fixed an issue with mark_compact_shape_dynamic + ### CUTLASS C++ * Further enhance Blackwell SM100 Attention kernels in [example 77](https://github.com/NVIDIA/cutlass/tree/main/examples/77_blackwell_fmha/). - Add softmax skip correction. - Fix a shared memory allocation bug where it needs to opt in maximum dynamics shared memory explicitly once it exceeds 48KB. - Fix a dead hang issue caused by early return warp. +* Add support through cmdline argument lists for `batch`, `no_verif`, `cluster_shape` and `cluster_shape_fallback` in [example 89](https://github.com/NVIDIA/cutlass/tree/main/examples/89_sm103_fp4_ultra_gemm/). * Add Ragged Contiguous Grouped gemm kernel in [example 92](https://github.com/NVIDIA/cutlass/tree/main/examples/92_blackwell_moe_gemm/). - This kernel uses a TMA 3D load to load the weights matrix and use the tensormap update method to load activations. +* Add 256x128 tile size support for Hopper SM90 deepgemm in [example 67](https://github.com/NVIDIA/cutlass/tree/main/examples/67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling/). + - Performance is optimized to align with Deepseek implementation. +* Simplification of API for MoE gemms. + - Instead of requiring users to call several cute utilities to set up the stride, API `moe_stride_utils` is introduced to help setup strides in the kernel. + - Instead of requiring users to set vectors like `problem_shapes_device` and `problem_shapes_hosts`, a new problem shape struct called `MoEProblemShape` is introduced which takes in max_m, max_n, max_k and counts vector as input and deduce problem shapes internally whenever required. +* Enable GEMM_K = 0 in grouped gemm. * Optimize group gemm kernels by enabling async TMA desc update. * Support Blackwell SM100 convolution stream-K kernel. - Unit tests: [fprop_streamK](https://github.com/NVIDIA/cutlass/tree/main/test/unit/conv/device_3x/fprop/sm100_conv3d_fprop_implicit_gemm_f16_f16_f16_tensorop_f16_streamk.cu), [dgrad_streamK](https://github.com/NVIDIA/cutlass/tree/main/test/unit/conv/device_3x/dgrad/sm100_conv3d_dgrad_implicit_gemm_f16_f16_f16_tensorop_f16_streamk.cu), [wgrad_streamK](https://github.com/NVIDIA/cutlass/tree/main/test/unit/conv/device_3x/wgrad/sm100_conv2d_wgrad_implicit_gemm_f16_f16_f16_tensorop_f16_streamk.cu). -* Add profiler support for Blackwell SM100 and SM120 blockscaled sparse kernels. +* Add Blackwell SM100 sparse gemm compressor unit tests. + - Unit tests: [compressor_fp16](https://github.com/NVIDIA/cutlass/tree/main/test/unit/transform/device/sm100_sparse_gemm_compressor_f16.cu). + - Add sub-bytes and runtime data type support in compressor unit test testbed. +* Add profiler support for: + - Blackwell SM100 and SM120 blockscaled sparse kernels. + - New MoE grouped gemm API. + - Blackwell SM100 cpasync kernel. * Fix some kernel issues: - Fix a race check issue of Blackwell SM103 kernels by adding missing elect one for prefetch barrier initialization. - Allow user to directly specify the number of stages for Hopper sm90 mixed input gemm. - Remove warnings caused by cuda vector type alignment setting in CUDA 13. - Remove problematic `cutlass::int8_t` and replace it with `int8_t`. + - Fix a few bugs in distributed gemm API and examples. + - Fix handling negative zero in sparse compressor. + - Add missing `wait_on_dependent_grids` for PDL use case. * Fix some profiler issues: - Add some missing reference kernels. + - Support VoidC reference kernels. - Add calculation of scale factor A and B in function `bytes_with_problem_shape` of block scaled profiler. + - Fix an issue when epilogue tile N is not divided by default subtile N. * Various improvements and fixes from the community and CUTLASS team. Thanks to everyone who submitted PRs! * Optimal code generation with CUDA toolkit versions 13.0U1. diff --git a/CMakeLists.txt b/CMakeLists.txt index 627d1713..b5340ba1 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1194,7 +1194,7 @@ if (CUTLASS_INSTALL_TESTS) install( FILES "${CMAKE_BINARY_DIR}/ctest/CTestTestfile.cmake" - DESTINATION "${CUTLASS_TEST_INSTALL_PREFIX}/" + DESTINATION "${CUTLASS_TEST_INSTALL_PREFIX}" ) endif() diff --git a/README.md b/README.md index 9e99803b..5470dd49 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ # CUTLASS 4.3.0 -_CUTLASS 4.3.0 - Oct 2025_ +_CUTLASS 4.3.0 - Nov 2025_ 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 @@ -46,20 +46,25 @@ To get started quickly - please refer : # What's New in CUTLASS 4.3 ## CuTe DSL +* New features: + - Supported Apache [TVM-FFI](https://tvm.apache.org/ffi/index.html) for further reduced host runtime overhead for JIT functions, better PyTorch and ML frameworks interopability + - Added fake tensor and stream to decouple compile jit function with "from_dlpack" flow. Now we no longer require users to have real tensor when compile jit function. + - Added FastDivmodDivisor with Python operator overloads, new APIs, Cute dialect integration, and optimized static tile scheduler performance for faster index mapping. + - Added l2 cache evict priority for tma related ops. Users could do fine-grain l2 cache control. * Debuggability improvements: - - Supported source location tracking for DSL APIs - - Supported dumping PTX and CUBIN code + - Supported source location tracking for DSL APIs (Allow tools like ``nsight`` profiling to correlate perf metrics with Python source code) + - Supported dumping PTX and CUBIN code: [Hello World Example](https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/notebooks/hello_world.ipynb) * More examples and notebooks to get started with CuTe DSL: - - [Kernel launch with Programmatic Dependent Launch](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/programmatic_dependent_launch.py) - - Improved performance of elementwise kernel (https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/ampere/elementwise_apply.py): + - Improved performance of [elementwise example](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/ampere/elementwise_apply.py): + Generalize code to handle list of input tensors + Generalize TV layout computation to handle different data types - - Demonstrate the new Pipeline APIs in [Blackwell SM100 persistent dense GEMM with static scheduling](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py): - + New Pipeline API `PipelineProducer` and `PipelineConsumer` to simplify code (no more explicit pipeline state management) - - Separate epilogue code for non-TMA and TMA implementation - + Note that the updates simplifies the codes but existing APIs still work and are supported - - [Basic Blackwell SM100 GEMM with decent performance](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py) - + Simple tutorial achieves 84% SOL performance with MNK 8K + - Improved [Blackwell SM100 persistent dense GEMM with static scheduling](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py): + + To demonstrate usage of new Pipeline APIs `PipelineProducer` and `PipelineConsumer` to simplify code without explicit pipeline state management (Exiting APIs are still maintained) + + Separated epilogue code for non-TMA and TMA implementation + - [Tutorial for Blackwell GEMM: Basic Blackwell SM100 GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/tutorial_gemm) + + [Baseline Blackwell GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py) achieves 84% SOL performance with MNK 8K + + More examples are coming for demo of optimization: `Baseline + X` + - [Tutorial for Async Pipeline API](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/async_pipeline.ipynb) - Reworked [elementwise add notebook](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/elementwise_add.ipynb) with more details and detailed explanation about TV layout + Updated implementation to handle general data type and multiple inputs + Updated explanation for TV layout in simpler language @@ -86,26 +91,48 @@ To get started quickly - please refer : - ``PipelineAsync`` - ``SmemAllocator`` - Fixed documentation for ``pipeline``, ``utils`` and ``cute.math`` + - Added overlapping accumulator optimization for block tile N = 256 case for better epilogue latency hiding 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). + - Fixed TensorSSA.__getitem__ indexing to match CuTe's indexing convention + - Fixed an issue with cutlass.max and cutlass.min + - Fixed an issue with mark_compact_shape_dynamic ## CUTLASS C++ * Further enhance Blackwell SM100 Attention kernels in [example 77](https://github.com/NVIDIA/cutlass/tree/main/examples/77_blackwell_fmha/). - Add softmax skip correction. - Fix a shared memory allocation bug where it needs to opt in maximum dynamics shared memory explicitly once it exceeds 48KB. - Fix a dead hang issue caused by early return warp. +* Add support through cmdline argument lists for `batch`, `no_verif`, `cluster_shape` and `cluster_shape_fallback` in [example 89](https://github.com/NVIDIA/cutlass/tree/main/examples/89_sm103_fp4_ultra_gemm/). * Add Ragged Contiguous Grouped gemm kernel in [example 92](https://github.com/NVIDIA/cutlass/tree/main/examples/92_blackwell_moe_gemm/). - This kernel uses a TMA 3D load to load the weights matrix and use the tensormap update method to load activations. +* Add 256x128 tile size support for Hopper SM90 deepgemm in [example 67](https://github.com/NVIDIA/cutlass/tree/main/examples/67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling/). + - Performance is optimized to align with Deepseek implementation. +* Simplification of API for MoE gemms. + - Instead of requiring users to call several cute utilities to set up the stride, API `moe_stride_utils` is introduced to help setup strides in the kernel. + - Instead of requiring users to set vectors like `problem_shapes_device` and `problem_shapes_hosts`, a new problem shape struct called `MoEProblemShape` is introduced which takes in max_m, max_n, max_k and counts vector as input and deduce problem shapes internally whenever required. +* Enable GEMM_K = 0 in grouped gemm. * Optimize group gemm kernels by enabling async TMA desc update. * Support Blackwell SM100 convolution stream-K kernel. - Unit tests: [fprop_streamK](https://github.com/NVIDIA/cutlass/tree/main/test/unit/conv/device_3x/fprop/sm100_conv3d_fprop_implicit_gemm_f16_f16_f16_tensorop_f16_streamk.cu), [dgrad_streamK](https://github.com/NVIDIA/cutlass/tree/main/test/unit/conv/device_3x/dgrad/sm100_conv3d_dgrad_implicit_gemm_f16_f16_f16_tensorop_f16_streamk.cu), [wgrad_streamK](https://github.com/NVIDIA/cutlass/tree/main/test/unit/conv/device_3x/wgrad/sm100_conv2d_wgrad_implicit_gemm_f16_f16_f16_tensorop_f16_streamk.cu). -* Add profiler support for Blackwell SM100 and SM120 blockscaled sparse kernels. +* Add Blackwell SM100 sparse gemm compressor unit tests. + - Unit tests: [compressor_fp16](https://github.com/NVIDIA/cutlass/tree/main/test/unit/transform/device/sm100_sparse_gemm_compressor_f16.cu). + - Add sub-bytes and runtime data type support in compressor unit test testbed. +* Add profiler support for: + - Blackwell SM100 and SM120 blockscaled sparse kernels. + - New MoE grouped gemm API. + - Blackwell SM100 cpasync kernel. * Fix some kernel issues: - Fix a race check issue of Blackwell SM103 kernels by adding missing elect one for prefetch barrier initialization. - Allow user to directly specify the number of stages for Hopper sm90 mixed input gemm. - Remove warnings caused by cuda vector type alignment setting in CUDA 13. - Remove problematic `cutlass::int8_t` and replace it with `int8_t`. + - Fix a few bugs in distributed gemm API and examples. + - Fix handling negative zero in sparse compressor. + - Add missing `wait_on_dependent_grids` for PDL use case. * Fix some profiler issues: - Add some missing reference kernels. + - Support VoidC reference kernels. - Add calculation of scale factor A and B in function `bytes_with_problem_shape` of block scaled profiler. + - Fix an issue when epilogue tile N is not divided by default subtile N. 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/bin2hex.cmake b/bin2hex.cmake index c03cdf78..284151fd 100644 --- a/bin2hex.cmake +++ b/bin2hex.cmake @@ -26,24 +26,55 @@ # 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 small utility function which generates a C-header from an input file -function(FILE_TO_C_STRING FILENAME VARIABLE_NAME OUTPUT_STRING ZERO_TERMINATED) - FILE(READ "${FILENAME}" HEX_INPUT HEX) - if (${ZERO_TERMINATED}) - string(APPEND HEX_INPUT "00") - endif() +# Option to enable faster raw string literal generation (non-MSVC compilers only) +# ON: Use raw string literals for faster compilation (default) +# OFF: Use hex array format for maximum compatibility +set(CUTLASS_BIN2HEX_FAST_MODE ON CACHE BOOL "Enable faster raw string generation for bin2hex") +# Helper function: Generate using raw string literal (fast mode) +function(cutlass_generate_raw_string_literal FILENAME VARIABLE_NAME OUTPUT_STRING) + # Read file as raw content + FILE(READ "${FILENAME}" RAW_CONTENT) + + # Unique delimiter that won't appear in the content + set(RAW_DELIM "abc1ae6e1e82e3de") + + # Create the raw string literal format: R"delimiter(...)delimiter" + set(RESULT "static char constexpr ${VARIABLE_NAME}[] = R\"${RAW_DELIM}(${RAW_CONTENT})${RAW_DELIM}\";\n") + + set(${OUTPUT_STRING} "${RESULT}" PARENT_SCOPE) +endfunction() + +# Helper function: Generate using hex array (compatible mode) +function(cutlass_generate_hex_array FILENAME VARIABLE_NAME OUTPUT_STRING) + # Read as hex and format + FILE(READ "${FILENAME}" HEX_INPUT HEX) + string(REGEX REPLACE "(....)" "\\1\n" HEX_OUTPUT ${HEX_INPUT}) string(REGEX REPLACE "([0-9a-f][0-9a-f])" "char(0x\\1)," HEX_OUTPUT ${HEX_OUTPUT}) + + set(RESULT "static char const ${VARIABLE_NAME}[] = {\n ${HEX_OUTPUT}\n};\n") + + set(${OUTPUT_STRING} "${RESULT}" PARENT_SCOPE) - set(HEX_OUTPUT "static char const ${VARIABLE_NAME}[] = {\n ${HEX_OUTPUT}\n};\n") +endfunction() - set(${OUTPUT_STRING} "${HEX_OUTPUT}" PARENT_SCOPE) +# Main function which generates a C-header from an input file +function(cutlass_file_to_c_string FILENAME VARIABLE_NAME OUTPUT_STRING) + # Use fast raw string literal mode only on non-MSVC compilers when fast mode is enabled + # MSVC has a 65,535 character limit for string literals, so always use hex array for safety + if(CUTLASS_BIN2HEX_FAST_MODE AND NOT MSVC) + cutlass_generate_raw_string_literal("${FILENAME}" "${VARIABLE_NAME}" RESULT) + else() + cutlass_generate_hex_array("${FILENAME}" "${VARIABLE_NAME}" RESULT) + endif() + + set(${OUTPUT_STRING} "${RESULT}" PARENT_SCOPE) endfunction() # message("Create header file for ${FILE_IN}") # message("Create header file for ${FILE_OUT}") -file_to_c_string(${FILE_IN} ${VARIABLE_NAME} OUTPUT_STRING ZERO_TERMINATED) +cutlass_file_to_c_string(${FILE_IN} ${VARIABLE_NAME} OUTPUT_STRING) set(RESULT "#pragma once\n") string(APPEND RESULT "namespace cutlass {\n") diff --git a/examples/52_hopper_gather_scatter_fusion/gather_gemm.hpp b/examples/52_hopper_gather_scatter_fusion/gather_gemm.hpp index 959e1e63..992836af 100644 --- a/examples/52_hopper_gather_scatter_fusion/gather_gemm.hpp +++ b/examples/52_hopper_gather_scatter_fusion/gather_gemm.hpp @@ -231,7 +231,7 @@ public: // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a. #if ! defined(__CUDA_ARCH_FEAT_SM90_ALL) if constexpr(size<0>(typename TiledMma::AtomShape_MNK{}) == 64) { - printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); return; } #endif diff --git a/examples/63_hopper_gemm_with_weight_prefetch/kernel/sm90_gemm_tma_warpspecialized_with_prefetch.hpp b/examples/63_hopper_gemm_with_weight_prefetch/kernel/sm90_gemm_tma_warpspecialized_with_prefetch.hpp index 73655ad2..ce746a6a 100644 --- a/examples/63_hopper_gemm_with_weight_prefetch/kernel/sm90_gemm_tma_warpspecialized_with_prefetch.hpp +++ b/examples/63_hopper_gemm_with_weight_prefetch/kernel/sm90_gemm_tma_warpspecialized_with_prefetch.hpp @@ -240,7 +240,7 @@ public: // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a. #if ! defined(ENABLE_SM90_KERNEL_LEVEL) - printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); #else enum class WarpGroupRole { diff --git a/examples/67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling/67_hopper_fp8_warp_specialized_gemm_with_groupwise_scaling.cu b/examples/67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling/67_hopper_fp8_warp_specialized_gemm_with_groupwise_scaling.cu index 19e012b0..e3a888e1 100644 --- a/examples/67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling/67_hopper_fp8_warp_specialized_gemm_with_groupwise_scaling.cu +++ b/examples/67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling/67_hopper_fp8_warp_specialized_gemm_with_groupwise_scaling.cu @@ -40,7 +40,7 @@ copies between thread blocks in a cluster. 3. This example uses the Warp Specialized kernel design (see /media/docs/efficient_gemm.md for details). 4. This example shows all important fusions used by FP8 gemm kernels, i.e., grouped scale factor along M for - A, blocked scale factor along K for A tensor, blocked scale factor for B tensor, the abs_max value of D tensor. + A, blocked scale factor along K for A tensor, blocked scale factor for B tensor. 5. A simple way to tune the CTA rasterization direction and swizzle pattern of Hopper kernels. Both the CTA rasterization direction and swizzle pattern impact cross-CTA locality of accesses. By tuning we can improve performance. @@ -71,6 +71,7 @@ #include "cutlass/util/host_tensor.h" #include "cutlass/util/packed_stride.hpp" #include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/reference/device/tensor_fill.h" #include "cutlass/util/reference/host/tensor_fill.h" #include "cutlass/util/reference/host/tensor_copy.h" #include "cutlass/util/reference/host/tensor_compare.h" @@ -100,30 +101,23 @@ using LayoutB = cutlass::layout::ColumnMajor; // L constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of B matrix in units of elements (up to 16 bytes) // C matrix configuration -using ElementC = float; // Element type for C and D matrix operands -using LayoutC = cutlass::layout::ColumnMajor; // Layout type for C and D matrix operands -constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes) +using ElementC = void; // Element type for C matrix operand +using LayoutC = cutlass::layout::ColumnMajor; // Layout type for C matrix operand +constexpr int AlignmentC = 1; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes) // D matrix configuration -using ElementD = ElementC; -using LayoutD = LayoutC; -constexpr int AlignmentD = AlignmentC; - -// Auxiliary matrix configuration and other fusion types -using ElementAux = ElementC; -using LayoutAux = LayoutC; -using ElementAmax = float; -using ElementBias = float; +using ElementD = cutlass::bfloat16_t; // Element type for D matrix operand +using LayoutD = cutlass::layout::RowMajor; // Layout type for D matrix operand +constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of D matrix in units of elements (up to 16 bytes) // Core kernel configurations using ElementAccumulator = float; // Element type for internal accumulation using ElementBlockScale = float; // Element type for blockscaling during accumulation using ElementCompute = float; // Element type for epilogue computation - -using ArchTag = cutlass::arch::Sm90; // Tag indicating the minimum SM that supports the intended feature -using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag -using TileShape = Shape<_128,_128,_128>; // Threadblock-level tile size -using ClusterShape = Shape<_1,_2,_1>; // Shape of the threadblocks in a cluster +using ArchTag = cutlass::arch::Sm90; // Tag indicating the minimum SM that supports the intended feature +using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag +using TileShape = Shape<_256,_128,_128>; // Threadblock-level tile size +using ClusterShape = Shape<_1,_2,_1>; // Shape of the threadblocks in a cluster constexpr int ScaleGranularityM = 1; constexpr int ScaleGranularityN = 128; @@ -134,14 +128,16 @@ constexpr int ScaleNsPerTile = size<1>(TileShape{}) / ScaleGranularityN; using ScaleConfig = cutlass::detail::Sm90BlockwiseScaleConfig; -using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); // Layout type for SFA matrix operand -using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); // Layout type for SFB matrix operand +using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); // Layout type for SFA matrix operand +using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); // Layout type for SFB matrix operand -using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8Blockwise; -using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative; -using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto; -using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltActAmaxAux< - LayoutAux, cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementC>; +using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8Blockwise; +using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative; + +using EpilogueTileType = Shape<_128,_128>; +using FusionOperation = cutlass::epilogue::fusion::Sm90EVT< + cutlass::epilogue::fusion::Sm90AccFetch + >; using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< ArchTag, OperatorClass, @@ -154,7 +150,7 @@ using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBui FusionOperation >::CollectiveOp; -using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< +using CollectiveMainloopWithBlockWiseScaling = typename cutlass::gemm::collective::CollectiveBuilder< ArchTag, OperatorClass, ElementA, cute::tuple, AlignmentA, ElementB, cute::tuple, AlignmentB, @@ -166,35 +162,22 @@ using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder KernelSchedule >::CollectiveOp; - using GemmKernel = cutlass::gemm::kernel::GemmUniversal< - Shape, - CollectiveMainloop, - CollectiveEpilogue, - cutlass::gemm::StreamKScheduler - >; + Shape, // Indicates ProblemShape + CollectiveMainloopWithBlockWiseScaling, + CollectiveEpilogue +>; using Gemm = cutlass::gemm::device::GemmUniversalAdapter; // Extract information from Gemm kernel. using EpilogueOutputOp = typename Gemm::EpilogueOutputOp; -using ElementScalar = typename EpilogueOutputOp::ElementScalar; -using ElementAmax = typename EpilogueOutputOp::ElementAmax; -using ActivationFunctor = typename EpilogueOutputOp::ActivationFn; +using ElementScalar = ElementCompute; using StrideA = typename Gemm::GemmKernel::StrideA; using StrideB = typename Gemm::GemmKernel::StrideB; using StrideC = typename Gemm::GemmKernel::StrideC; using StrideD = typename Gemm::GemmKernel::StrideD; -using StrideAux = StrideD; - -constexpr bool IsDFp8 = - cute::is_same_v or - cute::is_same_v; - -constexpr bool IsAuxFp8 = - cute::is_same_v or - cute::is_same_v; static_assert(cute::is_same_v, "ElementAccumulator and ElementBlockScale should be same datatype"); @@ -204,37 +187,22 @@ StrideA stride_A; StrideB stride_B; StrideC stride_C; StrideD stride_D; -StrideAux stride_aux; LayoutSFA layout_SFA; LayoutSFB layout_SFB; uint64_t seed; using LayoutScalar = cutlass::layout::PackedVectorLayout; - cutlass::HostTensor tensor_A; cutlass::HostTensor tensor_B; -cutlass::HostTensor tensor_C; cutlass::HostTensor tensor_D; cutlass::HostTensor blockscale_tensor_A; cutlass::HostTensor blockscale_tensor_B; cutlass::HostTensor tensor_ref_D; -cutlass::HostTensor tensor_aux; -cutlass::HostTensor tensor_ref_aux; -cutlass::HostTensor scalar_alpha; -cutlass::HostTensor scalar_beta; -cutlass::HostTensor scale_A; -cutlass::HostTensor scale_B; -cutlass::HostTensor scale_C; -cutlass::HostTensor scale_D; -cutlass::HostTensor scale_aux; -cutlass::HostTensor abs_max_D; -cutlass::HostTensor reference_abs_max_D; -cutlass::HostTensor abs_max_aux; -cutlass::HostTensor reference_abs_max_aux; #endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + ///////////////////////////////////////////////////////////////////////////////////////////////// /// Testbed utility types ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -272,7 +240,8 @@ template bool initialize_tensor( cutlass::TensorView view, cutlass::Distribution::Kind dist_kind, - uint64_t seed) { + uint64_t seed, + bool is_device_tensor = false) { if (dist_kind == cutlass::Distribution::Uniform) { @@ -294,22 +263,45 @@ bool initialize_tensor( scope_min = -8; } - cutlass::reference::host::TensorFillRandomUniform( - view, seed, scope_max, scope_min, bits_input); + if (is_device_tensor) { + using Real = typename cutlass::RealType::Type; + cutlass::reference::device::TensorFillRandomUniform( + view, seed, static_cast(scope_max), static_cast(scope_min), bits_input); + } else { + cutlass::reference::host::TensorFillRandomUniform( + view, seed, scope_max, scope_min, bits_input); + } + } else if (dist_kind == cutlass::Distribution::AllZeros) { - cutlass::reference::host::TensorFill(view); + if (is_device_tensor) { + cutlass::reference::device::TensorFill(view); + } else { + cutlass::reference::host::TensorFill(view); + } } else if (dist_kind == cutlass::Distribution::Identity) { - cutlass::reference::host::TensorFillIdentity(view); + if (is_device_tensor) { + cutlass::reference::device::TensorFillIdentity(view); + } else { + cutlass::reference::host::TensorFillIdentity(view); + } } else if (dist_kind == cutlass::Distribution::Gaussian) { - - cutlass::reference::host::TensorFillRandomGaussian(view, seed, 0, 0.5); + if (is_device_tensor) { + using Real = typename cutlass::RealType::Type; + cutlass::reference::device::TensorFillRandomGaussian(view, seed, static_cast(0), static_cast(0.5)); + } else { + cutlass::reference::host::TensorFillRandomGaussian(view, seed, 0, 0.5); + } } else if (dist_kind == cutlass::Distribution::Sequential) { - cutlass::reference::host::BlockFillSequential(view.data(), view.capacity()); + if (is_device_tensor) { + cutlass::reference::device::BlockFillSequential(view.data(), view.capacity()); + } else { + cutlass::reference::host::BlockFillSequential(view.data(), view.capacity()); + } } else { throw std::runtime_error("Not implementated."); @@ -323,7 +315,8 @@ template bool initialize_scale_tensor( cutlass::TensorView view, cutlass::Distribution::Kind dist_kind, - uint64_t seed) { + uint64_t seed, + bool is_device_tensor = false) { if (dist_kind == cutlass::Distribution::Uniform) { @@ -332,22 +325,45 @@ bool initialize_scale_tensor( scope_min = -1; scope_max = 1; - cutlass::reference::host::TensorFillRandomUniform( - view, seed, scope_max, scope_min); + if (is_device_tensor) { + using Real = typename cutlass::RealType::Type; + cutlass::reference::device::TensorFillRandomUniform( + view, seed, static_cast(scope_max), static_cast(scope_min)); + } else { + cutlass::reference::host::TensorFillRandomUniform( + view, seed, scope_max, scope_min); + } } else if (dist_kind == cutlass::Distribution::AllZeros) { - cutlass::reference::host::TensorFill(view); + if (is_device_tensor) { + cutlass::reference::device::TensorFill(view); + } else { + cutlass::reference::host::TensorFill(view); + } } else if (dist_kind == cutlass::Distribution::Identity) { - cutlass::reference::host::TensorFillIdentity(view); + if (is_device_tensor) { + cutlass::reference::device::TensorFillIdentity(view); + } else { + cutlass::reference::host::TensorFillIdentity(view); + } } else if (dist_kind == cutlass::Distribution::Gaussian) { - cutlass::reference::host::TensorFillRandomGaussian(view, seed, 0, 0.5); + if (is_device_tensor) { + using Real = typename cutlass::RealType::Type; + cutlass::reference::device::TensorFillRandomGaussian(view, seed, static_cast(0), static_cast(0.5)); + } else { + cutlass::reference::host::TensorFillRandomGaussian(view, seed, 0, 0.5); + } } else if (dist_kind == cutlass::Distribution::Sequential) { - cutlass::reference::host::BlockFillSequential(view.data(), view.capacity()); + if (is_device_tensor) { + cutlass::reference::device::BlockFillSequential(view.data(), view.capacity()); + } else { + cutlass::reference::host::BlockFillSequential(view.data(), view.capacity()); + } } else { throw std::runtime_error("Not implementated."); @@ -359,18 +375,19 @@ bool initialize_scale_tensor( /// Initialize operands to be used in the GEMM and reference GEMM void initialize(const Options &options) { - assert(options.m % ScaleGranularityM == 0); - assert(options.n % ScaleGranularityN == 0); - stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(options.m, options.k, options.l)); stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(options.n, options.k, options.l)); stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(options.m, options.n, options.l)); stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(options.m, options.n, options.l)); - stride_aux = stride_D; + + // Layout SFA and SFB represent logically broadcasting data in CuTe. + // E.g., if Layout SFA has shape ((ScaleGranularityM, M / ScaleGranularityM), (ScaleGraunularityK, K / ScaleGranularityK)) + // and strides ((0, 1), (0, M / ScaleGraunuarlityM)), then each collection of ScaleGranularityM x ScaleGranularityK + // indices in the tensor map to the same offset. + layout_SFA = ScaleConfig::tile_atom_to_shape_SFA(make_shape(options.m, options.n, options.k, options.l)); layout_SFB = ScaleConfig::tile_atom_to_shape_SFB(make_shape(options.m, options.n, options.k, options.l)); - auto a_coord = cutlass::make_Coord(options.m * options.l, options.k); auto c_coord = cutlass::make_Coord(options.m * options.l, options.n); auto b_coord = cutlass::make_Coord(options.k, options.n * options.l); @@ -378,94 +395,42 @@ void initialize(const Options &options) { auto groupscale_b_coord = cutlass::make_Coord(size(filter_zeros(layout_SFB))); tensor_A.resize(a_coord); - tensor_B.resize(b_coord); blockscale_tensor_A.resize(groupscale_a_coord); + tensor_B.resize(b_coord); blockscale_tensor_B.resize(groupscale_b_coord); - tensor_C.resize(c_coord); tensor_D.resize(c_coord); tensor_ref_D.resize(c_coord); - cutlass::Distribution::Kind dist_A = cutlass::Distribution::Uniform; - cutlass::Distribution::Kind dist_B = cutlass::Distribution::Uniform; - cutlass::Distribution::Kind dist_C = cutlass::Distribution::Uniform; - cutlass::Distribution::Kind dist_scaleA = cutlass::Distribution::Uniform; - cutlass::Distribution::Kind dist_scaleB = cutlass::Distribution::Uniform; + cutlass::Distribution::Kind dist_A = cutlass::Distribution::Gaussian; + cutlass::Distribution::Kind dist_B = cutlass::Distribution::Gaussian; + cutlass::Distribution::Kind dist_scaleA = cutlass::Distribution::Gaussian; + cutlass::Distribution::Kind dist_scaleB = cutlass::Distribution::Gaussian; - initialize_tensor(tensor_A.host_view(), dist_A, seed + 2022); - initialize_tensor(tensor_B.host_view(), dist_B, seed + 2023); - initialize_tensor(tensor_C.host_view(), dist_C, seed + 2024); - initialize_scale_tensor(blockscale_tensor_A.host_view(), dist_scaleA, seed + 2025); - initialize_scale_tensor(blockscale_tensor_B.host_view(), dist_scaleB, seed + 2026); + initialize_tensor(tensor_A.device_view(), dist_A, seed + 2022, true); + initialize_tensor(tensor_B.device_view(), dist_B, seed + 2023, true); + initialize_scale_tensor(blockscale_tensor_A.device_view(), dist_scaleA, seed + 2025, true); + initialize_scale_tensor(blockscale_tensor_B.device_view(), dist_scaleB, seed + 2026, true); + + tensor_A.sync_host(); + tensor_B.sync_host(); + tensor_D.sync_host(); + blockscale_tensor_A.sync_host(); + blockscale_tensor_B.sync_host(); #if 0 // Dump blockscaled tensors - std::cout << "blockscale_tensor_A: " << groupscale_a_coord << std::endl; + // Print block scaling tensors on the host side. + std::cout << "blockscale_tensor_A: " << blockscale_a_coord << std::endl; std::cout << blockscale_tensor_A.host_view() << "\n"; - std::cout << "blockscale_tensor_B: " << groupscale_b_coord << std::endl; + std::cout << "blockscale_tensor_B: " << blockscale_b_coord << std::endl; std::cout << blockscale_tensor_B.host_view() << "\n"; #endif - // Print group scaling tensors on the host side. - tensor_A.sync_device(); - tensor_B.sync_device(); - tensor_C.sync_device(); - tensor_D.sync_device(); - blockscale_tensor_A.sync_device(); - blockscale_tensor_B.sync_device(); - - if (options.save_aux) { - tensor_aux.resize(c_coord); - tensor_aux.sync_device(); - tensor_ref_aux.resize(c_coord); - } - - if (options.device_scale) { - scalar_alpha.resize(cutlass::make_Coord(1)); - scalar_beta.resize(cutlass::make_Coord(1)); - scale_A.resize(cutlass::make_Coord(1)); - scale_B.resize(cutlass::make_Coord(1)); - scale_C.resize(cutlass::make_Coord(1)); - scale_D.resize(cutlass::make_Coord(1)); - scale_aux.resize(cutlass::make_Coord(1)); - - cutlass::reference::host::TensorFill(scalar_alpha.host_view(), options.alpha); - cutlass::reference::host::TensorFill(scalar_beta.host_view(), options.beta); - cutlass::reference::host::TensorFill(scale_A.host_view(), options.scale_a); - cutlass::reference::host::TensorFill(scale_B.host_view(), options.scale_b); - cutlass::reference::host::TensorFill(scale_C.host_view(), options.scale_c); - cutlass::reference::host::TensorFill(scale_D.host_view(), options.scale_d); - cutlass::reference::host::TensorFill(scale_aux.host_view(), options.scale_aux); - - scalar_alpha.sync_device(); - scalar_beta.sync_device(); - scale_A.sync_device(); - scale_B.sync_device(); - scale_C.sync_device(); - scale_D.sync_device(); - scale_aux.sync_device(); - } - - if (IsDFp8 && options.save_amax) { - abs_max_D.resize(cutlass::make_Coord(1)); - initialize_tensor(abs_max_D.host_view(), cutlass::Distribution::AllZeros, 0); - abs_max_D.sync_device(); - reference_abs_max_D.resize(cutlass::make_Coord(1)); - initialize_tensor(reference_abs_max_D.host_view(), cutlass::Distribution::AllZeros, 0); - } - - if (IsAuxFp8 && options.save_aux && options.save_amax) { - abs_max_aux.resize(cutlass::make_Coord(1)); - initialize_tensor(abs_max_aux.host_view(), cutlass::Distribution::AllZeros, 0); - abs_max_aux.sync_device(); - reference_abs_max_aux.resize(cutlass::make_Coord(1)); - initialize_tensor(reference_abs_max_aux.host_view(), cutlass::Distribution::AllZeros, 0); - } } /// Populates a Gemm::Arguments structure from the given commandline options -template -GemmArguments args_from_options(const Options &options) +typename Gemm::Arguments args_from_options(const Options &options) { - GemmArguments arguments{ + typename Gemm::Arguments arguments{ cutlass::gemm::GemmUniversalMode::kGemm, {options.m, options.n, options.k, options.l}, {tensor_A.device_data(), @@ -479,44 +444,11 @@ GemmArguments args_from_options(const Options &options) }, { {}, // epilogue.thread - tensor_C.device_data(), stride_C, + nullptr, stride_C, tensor_D.device_data(), stride_D } }; - auto &fusion_args = arguments.epilogue.thread; - fusion_args.alpha = options.alpha; - fusion_args.beta = options.beta; - fusion_args.alpha_ptr = scalar_alpha.device_data(); - fusion_args.beta_ptr = scalar_beta.device_data(); - fusion_args.scale_a = options.scale_a; - fusion_args.scale_b = options.scale_b; - fusion_args.scale_c = options.scale_c; - fusion_args.scale_a_ptr = scale_A.device_data(); - fusion_args.scale_b_ptr = scale_B.device_data(); - fusion_args.scale_c_ptr = scale_C.device_data(); - - // ignored if tensor types are not fp8 - fusion_args.scale_d = options.scale_d; - fusion_args.scale_aux = options.scale_aux; - fusion_args.scale_d_ptr = scale_D.device_data(); - fusion_args.scale_aux_ptr = scale_aux.device_data(); - - // leaving/setting these as nullptr disables the fusion at runtime - fusion_args.bias_ptr = nullptr; - - if (options.save_aux) { - fusion_args.aux_ptr = tensor_aux.device_data(); - fusion_args.dAux = stride_aux; - if (options.save_amax) { - fusion_args.amax_aux_ptr = abs_max_aux.device_data(); - } - } - - if (options.save_amax) { - fusion_args.amax_D_ptr = abs_max_D.device_data(); - } - arguments.scheduler.raster_order = options.raster; // The tile scheduler will swizzle up to 8 and with the nearest multiple of 2 (i.e., 1, 2, 4, and 8) arguments.scheduler.max_swizzle_size = options.swizzle; @@ -524,7 +456,6 @@ GemmArguments args_from_options(const Options &options) return arguments; } -/// Don't know why the compiler does not like verify() being templated... bool verify(const Options &options) { // // Compute reference output @@ -543,24 +474,12 @@ bool verify(const Options &options) { stride_B ) ); - auto C = cute::make_tensor(tensor_C.host_data(), - cute::make_layout( - cute::make_shape(options.m, options.n, options.l), - stride_C - ) - ); auto D = cute::make_tensor(tensor_ref_D.host_data(), cute::make_layout( cute::make_shape(options.m, options.n, options.l), stride_D ) ); - auto Aux = cute::make_tensor(tensor_ref_aux.host_data(), - cute::make_layout( - cute::make_shape(options.m, options.n, options.l), - stride_aux - ) - ); auto SFA = cute::make_tensor(blockscale_tensor_A.host_data(), layout_SFA); auto SFB = cute::make_tensor(blockscale_tensor_B.host_data(), layout_SFB); @@ -569,8 +488,8 @@ bool verify(const Options &options) { cutlass::reference::host::GettBlockScalingMainloopParams< ElementAccumulator, - decltype(A), - decltype(SFA), + decltype(A), + decltype(SFA), decltype(B), decltype(SFB) > mainloop_params{A, SFA, B, SFB}; @@ -580,27 +499,18 @@ bool verify(const Options &options) { ElementScalar, ElementAccumulator, ElementCompute, - decltype(C), + unused_t, // C decltype(D), unused_t, // bias - decltype(Aux), + unused_t, // aux unused_t, // valpha - unused_t, // vbeta - ActivationFunctor + unused_t/*, // vbeta + ActivationFunctor*/ > epilogue_params; - epilogue_params.C = C; epilogue_params.D = D; - epilogue_params.Aux = Aux; - epilogue_params.alpha = options.alpha; - epilogue_params.beta = options.beta; - epilogue_params.scale_a = options.scale_a; - epilogue_params.scale_b = options.scale_b; - epilogue_params.scale_c = options.scale_c; - epilogue_params.scale_d = options.scale_d; - epilogue_params.scale_aux = options.scale_aux; - epilogue_params.abs_max_D = reference_abs_max_D.host_data(); - epilogue_params.abs_max_Aux = reference_abs_max_aux.host_data(); + epilogue_params.alpha = 1.0f; + epilogue_params.beta = 0.0f; // get reference result cutlass::reference::host::Gemm3x(mainloop_params, epilogue_params); @@ -608,7 +518,8 @@ bool verify(const Options &options) { // compare_reference bool passed = true; tensor_D.sync_host(); - passed &= cutlass::reference::host::TensorRelativelyEquals(tensor_D.host_view(), tensor_ref_D.host_view(), ElementAux(options.epsilon), ElementAux(options.non_zero_floor)); + passed &= cutlass::reference::host::TensorRelativelyEquals(tensor_D.host_view(), + tensor_ref_D.host_view(), ElementD(options.epsilon), ElementD(options.non_zero_floor)); double mse = cutlass::reference::host::TensorMSE(tensor_D.host_view(), tensor_ref_D.host_view()); double mre = cutlass::reference::host::TensorMRE(tensor_D.host_view(), tensor_ref_D.host_view()); double max_error = cutlass::reference::host::TensorGreatestError(tensor_D.host_view(), tensor_ref_D.host_view()); @@ -623,31 +534,13 @@ bool verify(const Options &options) { << "}" << std::endl; #endif - if (IsDFp8 && options.save_amax) { - abs_max_D.sync_host(); - std::cout << " Abs max D: " << abs_max_D.at(cutlass::make_Coord(0)) << ", reference: " << reference_abs_max_D.at(cutlass::make_Coord(0)) << std::endl; - passed &= cutlass::relatively_equal(abs_max_D.at(cutlass::make_Coord(0)), reference_abs_max_D.at(cutlass::make_Coord(0)), ElementScalar(options.epsilon), ElementScalar(options.non_zero_floor)); - } - - if (options.save_aux) { - tensor_aux.sync_host(); - passed &= cutlass::reference::host::TensorRelativelyEquals(tensor_aux.host_view(), tensor_ref_aux.host_view(), ElementAux(options.epsilon), ElementAux(options.non_zero_floor)); - mse = cutlass::reference::host::TensorMSE(tensor_aux.host_view(), tensor_ref_aux.host_view()); - mre = cutlass::reference::host::TensorMRE(tensor_aux.host_view(), tensor_ref_aux.host_view()); - max_error = cutlass::reference::host::TensorGreatestError(tensor_aux.host_view(), tensor_ref_aux.host_view()); - std::cout << " Aux MSE: " << mse << ", MRE: " << mre << ", greatest error: " << max_error << std::endl; - if (IsAuxFp8 && options.save_amax) { - abs_max_aux.sync_host(); - std::cout << " Abs max aux: " << abs_max_aux.at(cutlass::make_Coord(0)) << ", reference: " << reference_abs_max_aux.at(cutlass::make_Coord(0)) << std::endl; - passed &= cutlass::relatively_equal(abs_max_aux.at(cutlass::make_Coord(0)), reference_abs_max_aux.at(cutlass::make_Coord(0)), ElementScalar(options.epsilon), ElementScalar(options.non_zero_floor)); - } - } - return passed; } /// Execute a given example GEMM computation -int run(Options &options) { +template +int run(Options &options) +{ bool skip = false; std::cout << " Problem Size: " << options.m << 'x' << options.n << 'x' << options.k << 'x' << options.l << std::endl; @@ -680,7 +573,7 @@ int run(Options &options) { Gemm gemm; // Create a structure of gemm kernel arguments suitable for invoking an instance of Gemm - auto arguments = args_from_options(options); + auto arguments = args_from_options(options); // Using the arguments, query for extra workspace required for matrix multiplication computation size_t workspace_size = Gemm::get_workspace_size(arguments); @@ -715,7 +608,6 @@ int run(Options &options) { for (int iter = 0; iter < options.warmup + options.iterations; ++iter) { if (iter == options.warmup) timer.start(); - CUTLASS_CHECK(gemm.initialize(arguments, workspace.get())); CUTLASS_CHECK(gemm.run()); } timer.stop(); @@ -734,10 +626,10 @@ int run(Options &options) { raster = "Along M"; } + std::cout << " Problem Size: " << options.m << 'x' << options.n << 'x' << options.k << 'x' << options.l << std::endl; std::cout << " Rasterization: " << raster << " with a maximum CTA swizzle of " << options.swizzle << std::endl; std::cout << " Avg runtime: " << result.avg_runtime_ms << " ms" << std::endl; - std::cout << " GFLOPS: " << result.gflops << std::endl; - fflush(stdout); + std::cout << " GFLOPS: " << result.gflops << std::endl << std::endl << std::endl; } return result.passed; @@ -786,8 +678,7 @@ int main(int argc, char const **args) { // #if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) - bool passed = true; - passed = run(options); + bool passed = run(options); if (!passed) return -1; #endif diff --git a/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp b/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp index 742b507d..438ede78 100644 --- a/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp +++ b/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp @@ -1509,7 +1509,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized { CUTLASS_DEVICE void operator()(Params const& params, char* smem) { #if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)) - printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else int warp_idx = cutlass::canonical_warp_idx_sync(); auto role = warp_idx_to_role(warp_idx); diff --git a/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp b/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp index bf72843a..884c5751 100644 --- a/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp +++ b/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp @@ -1481,7 +1481,7 @@ struct Sm100FmhaBwdMlaKernelTmaWarpSpecialized { CUTLASS_DEVICE void operator()(Params const& params, char* smem) { #if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)) - printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else int warp_idx = cutlass::canonical_warp_idx_sync(); auto role = warp_idx_to_role(warp_idx); diff --git a/examples/77_blackwell_fmha/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp b/examples/77_blackwell_fmha/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp index a88b9a87..37c4418f 100644 --- a/examples/77_blackwell_fmha/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp +++ b/examples/77_blackwell_fmha/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp @@ -252,7 +252,7 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized { CUTLASS_DEVICE void operator()(const Params ¶ms, char* smem) { #if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)) - printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else TileScheduler tile_scheduler{params.tile_scheduler}; diff --git a/examples/77_blackwell_fmha/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp b/examples/77_blackwell_fmha/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp index d59b20ff..cd49b76f 100644 --- a/examples/77_blackwell_fmha/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp +++ b/examples/77_blackwell_fmha/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp @@ -248,7 +248,7 @@ struct Sm100FmhaGenKernelWarpspecialized { CUTLASS_DEVICE void operator()(const Params ¶ms, char* smem) { #if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)) - printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else TileScheduler tile_scheduler{params.tile_scheduler}; diff --git a/examples/77_blackwell_fmha/kernel/sm100_fmha_mla_tma_warpspecialized.hpp b/examples/77_blackwell_fmha/kernel/sm100_fmha_mla_tma_warpspecialized.hpp index 7f6a1bb9..aaea1ee8 100644 --- a/examples/77_blackwell_fmha/kernel/sm100_fmha_mla_tma_warpspecialized.hpp +++ b/examples/77_blackwell_fmha/kernel/sm100_fmha_mla_tma_warpspecialized.hpp @@ -508,7 +508,7 @@ struct Sm100FmhaMlaKernelTmaWarpspecialized { CUTLASS_DEVICE void operator()(Params const& params, char* smem_raw) { #if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)) - printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else TileScheduler tile_scheduler(params.tile_scheduler); diff --git a/examples/88_hopper_fmha/kernel/fmha_kernel_tma.hpp b/examples/88_hopper_fmha/kernel/fmha_kernel_tma.hpp index eeee7fa2..dc08fc77 100644 --- a/examples/88_hopper_fmha/kernel/fmha_kernel_tma.hpp +++ b/examples/88_hopper_fmha/kernel/fmha_kernel_tma.hpp @@ -138,7 +138,7 @@ struct FmhaKernelTma { CUTLASS_DEVICE void operator()(const Params ¶ms, char* smem) { #if ! defined(CUTLASS_ARCH_MMA_SM90A_ENABLED) - printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else TileScheduler tile_scheduler{params.tile_scheduler}; diff --git a/examples/88_hopper_fmha/kernel/fmha_kernel_tma_warpspecialized.hpp b/examples/88_hopper_fmha/kernel/fmha_kernel_tma_warpspecialized.hpp index 4b96d711..7a2a6660 100644 --- a/examples/88_hopper_fmha/kernel/fmha_kernel_tma_warpspecialized.hpp +++ b/examples/88_hopper_fmha/kernel/fmha_kernel_tma_warpspecialized.hpp @@ -161,7 +161,7 @@ struct FmhaKernelTmaWarpSpecialized { CUTLASS_DEVICE void operator()(const Params ¶ms, char* smem) { #if ! defined(CUTLASS_ARCH_MMA_SM90A_ENABLED) - printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else enum class WarpGroupRole { diff --git a/examples/89_sm103_fp4_ultra_gemm/89_sm103_fp4_ultra_gemm.cu b/examples/89_sm103_fp4_ultra_gemm/89_sm103_fp4_ultra_gemm.cu index d6d7879c..9946e878 100644 --- a/examples/89_sm103_fp4_ultra_gemm/89_sm103_fp4_ultra_gemm.cu +++ b/examples/89_sm103_fp4_ultra_gemm/89_sm103_fp4_ultra_gemm.cu @@ -107,39 +107,62 @@ using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; // O // using ElementD = cutlass::float_e2m1_t; // Enable for SF Output // Element type for D matrix operands // Kernel Perf config -using MmaTileShape = cute::Shape>; // MMA's tile size -using ClusterShape = cute::Shape; // Shape of the threadblocks in a cluster +using MmaTileShape1Sm = cute::Shape>;// 1SM MMA's tile size +using MmaTileShape2Sm = cute::Shape>;// 2SM MMA's tile size +using ClusterShape = cute::Shape; // Cluster shape -// Epilogue fusion operator -using EpilogueFusionOp = cutlass::epilogue::fusion::LinearCombination; - -using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< +using CollectiveEpilogue1Sm = typename cutlass::epilogue::collective::CollectiveBuilder< ArchTag, OperatorClass, - MmaTileShape, ClusterShape, + MmaTileShape1Sm, ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, ElementAccumulator, ElementC, LayoutCTag, AlignmentC, ElementD, LayoutDTag, AlignmentD, - cutlass::epilogue::NoSmemWarpSpecialized1Sm, - EpilogueFusionOp + cutlass::epilogue::NoSmemWarpSpecialized1Sm >::CollectiveOp; -using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< +using CollectiveEpilogue2Sm = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, + MmaTileShape2Sm, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, LayoutCTag, AlignmentC, + ElementD, LayoutDTag, AlignmentD, + cutlass::epilogue::NoSmemWarpSpecialized2Sm + >::CollectiveOp; + +using CollectiveMainloop1Sm = typename cutlass::gemm::collective::CollectiveBuilder< ArchTag, OperatorClass, cute::tuple, LayoutATag, AlignmentA, cute::tuple, LayoutBTag, AlignmentB, ElementAccumulator, - MmaTileShape, ClusterShape, - cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, - cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103 // Kernel schedule policy. Auto or using targeted scheduling policy + MmaTileShape1Sm, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue1Sm::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103 // Kernel schedule policy. Auto or using targeted scheduling policy + >::CollectiveOp; +using CollectiveMainloop2Sm = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + cute::tuple, LayoutATag, AlignmentA, + cute::tuple, LayoutBTag, AlignmentB, + ElementAccumulator, + MmaTileShape2Sm, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue2Sm::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103 // Kernel schedule policy. Auto or using targeted scheduling policy >::CollectiveOp; -using GemmKernel = cutlass::gemm::kernel::GemmUniversal< +using GemmKernel1Sm = cutlass::gemm::kernel::GemmUniversal< Shape, // Indicates ProblemShape - CollectiveMainloop, - CollectiveEpilogue>; + CollectiveMainloop1Sm, + CollectiveEpilogue1Sm>; -using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +using Gemm1Sm = cutlass::gemm::device::GemmUniversalAdapter; +using Gemm = Gemm1Sm; +using GemmKernel2Sm = cutlass::gemm::kernel::GemmUniversal< + Shape, // Indicates ProblemShape + CollectiveMainloop2Sm, + CollectiveEpilogue2Sm>; + +using Gemm2Sm = cutlass::gemm::device::GemmUniversalAdapter; // Reference device GEMM implementation type using StrideA = typename Gemm::GemmKernel::StrideA; @@ -202,6 +225,12 @@ struct Options { int m, n, k; int swizzle = 0; + dim3 cluster_shape = dim3(2,1,1); + dim3 cluster_shape_fallback = dim3(2,1,1); + + bool verification = true; + int batch = 1; + Options(): help(false), m(1024), n(1024), k(1024), @@ -226,6 +255,14 @@ struct Options { cmd.get_cmd_line_argument("beta", beta, 0.f); cmd.get_cmd_line_argument("iterations", iterations); cmd.get_cmd_line_argument("swizzle", swizzle); + cmd.get_cmd_line_argument("cluster_m", cluster_shape.x); + cmd.get_cmd_line_argument("cluster_n", cluster_shape.y); + cmd.get_cmd_line_argument("cluster_fallback_m", cluster_shape_fallback.x); + cmd.get_cmd_line_argument("cluster_fallback_n", cluster_shape_fallback.y); + if (cmd.check_cmd_line_flag("no_verif")) { + verification = false; + } + cmd.get_cmd_line_argument("batch", batch); } /// Prints the usage statement. @@ -240,11 +277,19 @@ struct Options { << " --k= Sets the K extent of the GEMM\n" << " --alpha= Epilogue scalar alpha\n" << " --beta= Epilogue scalar beta\n" + << " --cluster_m= Preferred cluster X dimension (input only)\n" + << " --cluster_n= Preferred cluster Y dimension (input only)\n" + << " --cluster_fallback_m= Fallback cluster X dimension (input only)\n" + << " --cluster_fallback_n= Fallback cluster Y dimension (input only)\n" << " --swizzle= Cluster rasterization swizzle\n" + << " --batch= Number of batches (L dimension)\n" + << " --no_verif Do not run host-side verification\n" << " --iterations= Number of profiling iterations to perform.\n\n"; out << "\n\nExamples:\n\n" - << "$ " << "./examples/89_sm103_fp4_ultra_gemm/89_sm103_fp4_ultra_gemm" << " --m=1024 --n=512 --k=1024 --alpha=2 --beta=0.707 \n\n"; + << "$ " << "./examples/89_sm103_fp4_ultra_gemm/89_sm103_fp4_ultra_gemm" + << " --m=1024 --n=512 --k=1024 --alpha=2 --beta=0.707" + << " --cluster_m=4 --cluster_n=4 --cluster_fallback_m=2 --cluster_fallback_n=1\n\n"; return out; } @@ -252,8 +297,8 @@ struct Options { /// Compute performance in GFLOP/s double gflops(double runtime_s) const { - // Two flops per multiply-add - uint64_t flop = uint64_t(2) * m * n * k; + // Two flops per multiply-add, times batch (L dimension) + uint64_t flop = uint64_t(2) * uint64_t(m) * uint64_t(n) * uint64_t(k) * uint64_t(batch); double gflop = double(flop) / double(1.0e9); return gflop / runtime_s; } @@ -328,17 +373,17 @@ void initialize(const Options &options) { // For SFA and SFB tensors layouts using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; - stride_A = cutlass::make_cute_packed_stride(StrideA{}, {options.m, options.k, 1}); - stride_B = cutlass::make_cute_packed_stride(StrideB{}, {options.n, options.k, 1}); - stride_C = cutlass::make_cute_packed_stride(StrideC{}, {options.m, options.n, 1}); - stride_D = cutlass::make_cute_packed_stride(StrideD{}, {options.m, options.n, 1}); + stride_A = cutlass::make_cute_packed_stride(StrideA{}, {options.m, options.k, options.batch}); + stride_B = cutlass::make_cute_packed_stride(StrideB{}, {options.n, options.k, options.batch}); + stride_C = cutlass::make_cute_packed_stride(StrideC{}, {options.m, options.n, options.batch}); + stride_D = cutlass::make_cute_packed_stride(StrideD{}, {options.m, options.n, options.batch}); - layout_A = make_layout(make_shape(options.m, options.k, 1), stride_A); - layout_B = make_layout(make_shape(options.n, options.k, 1), stride_B); - layout_C = make_layout(make_shape(options.m, options.n, 1), stride_C); - layout_D = make_layout(make_shape(options.m, options.n, 1), stride_D); - layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(options.m, options.n, options.k, 1)); - layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(options.m, options.n, options.k, 1)); + layout_A = make_layout(make_shape(options.m, options.k, options.batch), stride_A); + layout_B = make_layout(make_shape(options.n, options.k, options.batch), stride_B); + layout_C = make_layout(make_shape(options.m, options.n, options.batch), stride_C); + layout_D = make_layout(make_shape(options.m, options.n, options.batch), stride_D); + layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(options.m, options.n, options.k, options.batch)); + layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(options.m, options.n, options.k, options.batch)); block_A.reset(cutlass::make_Coord(size(layout_A))); block_B.reset(cutlass::make_Coord(size(layout_B))); @@ -362,11 +407,12 @@ void initialize(const Options &options) { } // Populates a Gemm::Arguments structure from the given commandline options +template typename Gemm::Arguments args_from_options(const Options &options) { typename Gemm::Arguments arguments { cutlass::gemm::GemmUniversalMode::kGemm, - {options.m, options.n, options.k, 1}, + {options.m, options.n, options.k, options.batch}, { // Mainloop arguments block_A.device_data(), stride_A, block_B.device_data(), stride_B, @@ -381,6 +427,8 @@ typename Gemm::Arguments args_from_options(const Options &options) }; arguments.scheduler.max_swizzle_size = options.swizzle; + arguments.hw_info.cluster_shape = options.cluster_shape; + arguments.hw_info.cluster_shape_fallback = options.cluster_shape_fallback; return arguments; } @@ -432,7 +480,7 @@ int run(Options &options) Gemm gemm; // Create a structure of gemm kernel arguments suitable for invoking an instance of Gemm - auto arguments = args_from_options(options); + auto arguments = args_from_options(options); // Using the arguments, query for extra workspace required for matrix multiplication computation size_t workspace_size = Gemm::get_workspace_size(arguments); @@ -461,12 +509,15 @@ int run(Options &options) // Check if output from CUTLASS kernel and reference kernel are equal or not Result result; - result.passed = verify(options); - - std::cout << " Disposition: " << (result.passed ? "Passed" : "Failed") << std::endl; - - if (!result.passed) { - exit(-1); + if (options.verification) { + result.passed = verify(options); + std::cout << " Disposition: " << (result.passed ? "Passed" : "Failed") << std::endl; + if (!result.passed) { + exit(-1); + } + } else { + std::cout << " Disposition: Skipped verification" << std::endl; + result.passed = true; } // Run profiling loop @@ -486,7 +537,7 @@ int run(Options &options) result.gflops = options.gflops(result.avg_runtime_ms / 1000.0); - std::cout << " Problem Size: " << options.m << 'x' << options.n << 'x' << options.k << std::endl; + std::cout << " Problem Size: " << options.m << 'x' << options.n << 'x' << options.k << " (batch: " << options.batch << ")" << std::endl; std::cout << " Avg runtime: " << result.avg_runtime_ms << " ms" << std::endl; std::cout << " GFLOPS: " << result.gflops << std::endl; } @@ -536,7 +587,10 @@ int main(int argc, char const **args) { // Evaluate CUTLASS kernels // #if defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED) - run(options); + std::cout << "Running kernel with 1SM MMA config:" << std::endl; + run(options); + std::cout << "Running kernel with 2SM MMA config:" << std::endl; + run(options); #endif // defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED) return 0; diff --git a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_fp4_grouped.cu b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_fp4_grouped.cu index aef683cb..7f24e92f 100644 --- a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_fp4_grouped.cu +++ b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_fp4_grouped.cu @@ -39,7 +39,7 @@ Usage: $ ./examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_fp4_grouped - --m=28672 --n=4 --k=4096 --l=8 --benchmark=benchmark.txt + --m=7168 --n=128 --k=512 --group=8 --benchmark=benchmark.txt */ @@ -112,7 +112,7 @@ struct Options { bool error; bool verification; - int m, n, k, l; + int m, n, k, groups; int iterations; @@ -122,7 +122,7 @@ struct Options { help(false), error(false), verification(true), - m(2048), n(2048), k(2048), l(1), + m(2048), n(2048), k(2048), groups(1), iterations(10) { } @@ -138,7 +138,7 @@ struct Options { cmd.get_cmd_line_argument("m", m, 2048); cmd.get_cmd_line_argument("n", n, 2048); cmd.get_cmd_line_argument("k", k, 2048); - cmd.get_cmd_line_argument("l", l, 1); + cmd.get_cmd_line_argument("groups", groups, 1); cmd.get_cmd_line_argument("iterations", iterations, 10); cmd.get_cmd_line_argument("benchmark", benchmark_path); @@ -158,7 +158,7 @@ struct Options { << " --m= Sets the M extent of the GEMM\n" << " --n= Sets the N extent of the GEMM\n" << " --k= Sets the K extent of the GEMM\n" - << " --l= Sets the L extent (batch count) of the GEMM\n" + << " --groups= Sets the groups extent (batch count) of the GEMM\n" << " --iterations= Set the number of profiling iterations to perform\n" << " --benchmark= Executes a benchmark problem size\n" << " --no_verif Do not run verification kernels\n"; @@ -283,9 +283,7 @@ struct ExampleRunner { MainloopScheduleType >::CollectiveOp; - using ProblemShapeGroup = cutlass::gemm::GroupProblemShape>; // per group - using ProblemShapeMax = Shape; // max - using ProblemShape = cutlass::gemm::MoEProblemShape; + using ProblemShape = cutlass::gemm::MoEProblemShape>; // per group using GemmKernel = cutlass::gemm::kernel::GemmUniversal< ProblemShape, @@ -340,7 +338,11 @@ struct ExampleRunner { cutlass::HostTensor block_reference_SFD; cutlass::HostTensor block_Normconst; - cutlass::DeviceAllocation problem_sizes; + cutlass::DeviceAllocation tokens_per_expert; + + std::vector problem_sizes_host; + std::vector tokens_per_expert_host; + // // Methods @@ -397,10 +399,11 @@ struct ExampleRunner { // Comparison block_D.sync_host(); - auto [maxM, maxN, maxK, L] = problem_size.max_problem_shape; - for (int i = 0; i < problem_size.problem_shape.num_groups; i++) { - auto problem = problem_size.problem_shape.get_host_problem_shape(i); + auto [maxM, maxN, maxK] = problem_size.get_host_problem_shape(0); //gets max problem shape + for (int i = 0; i < problem_size.num_groups; i++) { + auto problem = problem_sizes_host.at(i); auto [M, N, K] = problem; + printf("group [%d] : M = %d, N = %d, K = %d\n", i, M, N, K); // assume all M == maxM auto refD_view = block_reference_D.host_view().subview(cutlass::make_Coord(M * N), cutlass::make_Coord(i * maxN * maxM)); @@ -414,7 +417,7 @@ struct ExampleRunner { /// Initialize operands to be used in the GEMM and reference GEMM void initialize(ProblemShape const& problem_size) { - auto problem_shape_MNKL = cute::append<4>(problem_size.max_problem_shape, 1); + auto problem_shape_MNKL = cute::append<4>(problem_size.get_host_problem_shape(0), problem_size.groups()); auto [M, N, K, L] = problem_shape_MNKL; // For SFA and SFB tensors layouts @@ -422,17 +425,11 @@ struct ExampleRunner { // For SFD tensor layout using Sm1xxBlockScaledOutputConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; - // printf("\nStrideC = "); - // print(StrideC{}); - stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, L}); stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, L}); stride_C = cutlass::make_cute_packed_stride(StrideC{}, {M, N, L}); stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, L}); - // printf("\nstride_C = "); - // print(stride_C); - layout_A = make_layout(make_shape(M, K, L), stride_A); layout_B = make_layout(make_shape(N, K, L), stride_B); layout_C = make_layout(make_shape(M, N, L), stride_C); @@ -441,16 +438,6 @@ struct ExampleRunner { layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, L)); layout_SFD = SfdOutputCfg::tile_atom_to_shape_SFD(cute::make_shape(M, N, K, L)); - // printf("\nlayout_A = "); - // print(layout_A); - // printf("\nlayout_B = "); - // print(layout_B); - // printf("\nlayout_C = "); - // print(layout_C); - - // printf("\nsize(layout_A)=%lld", (long long)size(layout_A)); - // printf("\n"); - block_A.reset(cutlass::make_Coord(size(layout_A))); block_B.reset(cutlass::make_Coord(size(layout_B))); block_C.reset(cutlass::make_Coord(size(layout_C))); @@ -481,12 +468,12 @@ struct ExampleRunner { } /// Load a benchmark - std::vector benchmark_problems(std::string const& benchmark_path) { - std::vector problem_sizes_host; + void benchmark_problems(std::string const& benchmark_path) { std::ifstream file(benchmark_path); if (!file.good()) { - return {}; + std::cout << "Issues with benchmark file." << std::endl; + return; } while (file.good()) { @@ -511,34 +498,68 @@ struct ExampleRunner { problem_sizes_host.push_back({extent.m(), extent.n(), extent.k()}); } - return problem_sizes_host; + return; + } + + void benchmark_tokens_per_expert(std::string const& benchmark_path) { + + std::ifstream file(benchmark_path); + if (!file.good()) { + return; + } + + while (file.good()) { + + int idx = -1; + std::string extent_str; + + file >> idx >> extent_str; + + if (idx < 0 || extent_str.empty()) { + break; + } + + cutlass::gemm::GemmCoord extent; + std::vector tokens; + + cutlass::CommandLine::tokenize(tokens, extent_str, 'x'); + + for (int i = 0; i < int(tokens.size()); ++i) { + extent.at(i) = std::atoi(tokens.at(i).c_str()); + } + tokens_per_expert_host.push_back(extent.n()); + } + + return; } bool run(Options const& options, cutlass::KernelHardwareInfo const& hw_info) { - auto problem_sizes_host = benchmark_problems(options.benchmark_path); + + benchmark_problems(options.benchmark_path); if (problem_sizes_host.empty()) { return false; } - problem_sizes.reset(problem_sizes_host.size()); - problem_sizes.copy_from_host(problem_sizes_host.data()); + benchmark_tokens_per_expert(options.benchmark_path); + if (tokens_per_expert_host.empty()) { + return false; + } - ProblemShape problem_size; - problem_size.max_problem_shape = ProblemShapeMax{options.m, options.n, options.k, options.l}; - problem_size.problem_shape.num_groups = problem_sizes_host.size(); - problem_size.problem_shape.problem_shapes = problem_sizes.get(); - problem_size.problem_shape.host_problem_shapes = problem_sizes_host.data(); + tokens_per_expert.reset(tokens_per_expert_host.size()); + tokens_per_expert.copy_from_host(tokens_per_expert_host.data()); + + ProblemShape problem_size {options.m, options.n, options.k, options.groups, tokens_per_expert.get(), tokens_per_expert_host.data()}; initialize(problem_size); - typename Gemm::Arguments arguments { + typename Gemm::Arguments arguments{ cutlass::gemm::GemmUniversalMode::kGrouped, problem_size, { // Mainloop arguments - block_A.device_data(), stride_A, - block_B.device_data(), stride_B, - block_SFA.device_data(), layout_SFA, - block_SFB.device_data(), layout_SFB + block_A.device_data(), + block_B.device_data(), + block_SFA.device_data(), + block_SFB.device_data() }, { // Epilogue arguments {}, @@ -624,7 +645,7 @@ struct ExampleRunner { // Compute average setup and runtime and FLOPs. float elapsed_ms = timer.elapsed_millis(); double avg_runtime_ms = double(elapsed_ms) / double(options.iterations); - double flops = double(int64_t(2) * options.m * options.n * options.k * options.l) / (avg_runtime_ms / 1000.0); + double flops = double(int64_t(2) * options.m * options.n * options.k * options.groups) / (avg_runtime_ms / 1000.0); std::cout << " Avg runtime : " << avg_runtime_ms << " ms" << std::endl; std::cout << " TFLOPS : " << flops / 1e12 << std::endl; diff --git a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_fp4_regular.cu b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_fp4_regular.cu index e129d07e..9ede1ef2 100644 --- a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_fp4_regular.cu +++ b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_fp4_regular.cu @@ -423,17 +423,11 @@ struct ExampleRunner { // For SFD tensor layout using Sm1xxBlockScaledOutputConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; - // printf("\nStrideC = "); - // print(StrideC{}); - stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, L}); stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, L}); stride_C = cutlass::make_cute_packed_stride(StrideC{}, {M, N, L}); stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, L}); - // printf("\nstride_C = "); - // print(stride_C); - layout_A = make_layout(make_shape(M, K, L), stride_A); layout_B = make_layout(make_shape(N, K, L), stride_B); layout_C = make_layout(make_shape(M, N, L), stride_C); @@ -442,16 +436,6 @@ struct ExampleRunner { layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, L)); layout_SFD = SfdOutputCfg::tile_atom_to_shape_SFD(cute::make_shape(M, N, K, L)); - // printf("\nlayout_A = "); - // print(layout_A); - // printf("\nlayout_B = "); - // print(layout_B); - // printf("\nlayout_C = "); - // print(layout_C); - - // printf("\nsize(layout_A)=%lld", (long long)size(layout_A)); - // printf("\n"); - block_A.reset(cutlass::make_Coord(size(layout_A))); block_B.reset(cutlass::make_Coord(size(layout_B))); block_C.reset(cutlass::make_Coord(size(layout_C))); @@ -481,27 +465,53 @@ struct ExampleRunner { block_Normconst.sync_device(); } + /// Populates a Gemm::Arguments structure from the given commandline options + typename Gemm::Arguments args_from_options(ProblemShapeType problem_size, cutlass::KernelHardwareInfo const& hw_info) + { + if constexpr (std::is_same_v){ + return typename Gemm::Arguments { + cutlass::gemm::GemmUniversalMode::kGemm, + problem_size, + { // Mainloop arguments + block_A.device_data(), + block_B.device_data(), + block_SFA.device_data(), + block_SFB.device_data() + }, + { // Epilogue arguments + {}, + block_C.device_data(), stride_C, + block_D.device_data(), stride_D + }, + hw_info + }; + } + else { + return typename Gemm::Arguments { + cutlass::gemm::GemmUniversalMode::kGemm, + problem_size, + { // Mainloop arguments + block_A.device_data(), stride_A, + block_B.device_data(), stride_B, + block_SFA.device_data(), layout_SFA, + block_SFB.device_data(), layout_SFB + }, + { // Epilogue arguments + {}, + block_C.device_data(), stride_C, + block_D.device_data(), stride_D + }, + hw_info + }; + } + } + bool run(Options const& options, cutlass::KernelHardwareInfo const& hw_info) { ProblemShapeType problem_size = ProblemShapeType{options.m, options.n, options.k, options.l}; initialize(problem_size); - typename Gemm::Arguments arguments { - cutlass::gemm::GemmUniversalMode::kGemm, - problem_size, - { // Mainloop arguments - block_A.device_data(), stride_A, - block_B.device_data(), stride_B, - block_SFA.device_data(), layout_SFA, - block_SFB.device_data(), layout_SFB - }, - { // Epilogue arguments - {}, - block_C.device_data(), stride_C, - block_D.device_data(), stride_D - }, - hw_info - }; + typename Gemm::Arguments arguments = args_from_options(problem_size, hw_info); if constexpr (IsBlockScaleSupported) { arguments.epilogue.thread.block_scale_factor_ptr = block_SFD.device_data(); diff --git a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_grouped.cu b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_grouped.cu index 8e5325af..cd7080d6 100644 --- a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_grouped.cu +++ b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_grouped.cu @@ -39,7 +39,7 @@ Usage: $ ./examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_grouped - --m=28672 --n=4 --k=4096 --l=8 --benchmark=benchmark.txt + --m=7168 --n=128 --k=512 --groups=8 --benchmark=benchmark.txt */ @@ -82,7 +82,7 @@ struct Options { bool error; bool verification; - int m, n, k, l; + int m, n, k, groups; int iterations; @@ -92,7 +92,7 @@ struct Options { help(false), error(false), verification(true), - m(2048), n(2048), k(2048), l(1), + m(2048), n(2048), k(2048), groups(1), iterations(10) { } @@ -108,7 +108,7 @@ struct Options { cmd.get_cmd_line_argument("m", m, 2048); cmd.get_cmd_line_argument("n", n, 2048); cmd.get_cmd_line_argument("k", k, 2048); - cmd.get_cmd_line_argument("l", l, 1); + cmd.get_cmd_line_argument("groups", groups, 1); cmd.get_cmd_line_argument("iterations", iterations, 10); cmd.get_cmd_line_argument("benchmark", benchmark_path); @@ -128,7 +128,7 @@ struct Options { << " --m= Sets the M extent of the GEMM\n" << " --n= Sets the N extent of the GEMM\n" << " --k= Sets the K extent of the GEMM\n" - << " --l= Sets the L extent (batch count) of the GEMM\n" + << " --groups= Sets the groups extent (batch count) of the GEMM\n" << " --iterations= Set the number of profiling iterations to perform\n" << " --benchmark= Executes a benchmark problem size\n" << " --no_verif Do not run verification kernels\n"; @@ -219,9 +219,7 @@ struct ExampleRunner { MainloopScheduleType >::CollectiveOp; - using ProblemShapeGroup = cutlass::gemm::GroupProblemShape>; // per group - using ProblemShapeMax = Shape; // max - using ProblemShape = cutlass::gemm::MoEProblemShape; + using ProblemShape = cutlass::gemm::MoEProblemShape>; // per group using GemmKernel = cutlass::gemm::kernel::GemmUniversal< ProblemShape, @@ -232,8 +230,6 @@ struct ExampleRunner { using Gemm = cutlass::gemm::device::GemmUniversalAdapter; - // using ProblemShapeType = typename Gemm::GemmKernel::ProblemShape; - using StrideA = typename Gemm::GemmKernel::StrideA; using StrideB = typename Gemm::GemmKernel::StrideB; using StrideC = typename Gemm::GemmKernel::StrideC; @@ -249,8 +245,6 @@ struct ExampleRunner { // /// Initialization - StrideA stride_A; - StrideB stride_B; StrideC stride_C; StrideD stride_D; uint64_t seed = 0; @@ -261,7 +255,10 @@ struct ExampleRunner { cutlass::DeviceAllocation block_D; cutlass::DeviceAllocation block_ref_D; - cutlass::DeviceAllocation problem_sizes; + cutlass::DeviceAllocation tokens_per_expert; + + std::vector problem_sizes_host; + std::vector tokens_per_expert_host; // @@ -269,10 +266,11 @@ struct ExampleRunner { // bool verify(ProblemShape const& problem_size, float alpha, float beta) { - auto [maxM, maxN, maxK, L] = problem_size.max_problem_shape; - for (int i = 0; i < problem_size.problem_shape.num_groups; i++) { - auto problem = problem_size.problem_shape.get_host_problem_shape(i); + auto [maxM, maxN, maxK] = problem_size.get_host_problem_shape(0); //gets max problem shape + for (int i = 0; i < problem_size.num_groups; i++) { + auto problem = problem_sizes_host.at(i); auto [M, N, K] = problem; + printf("group [%d] : M = %d, N = %d, K = %d\n", i, M, N, K); cutlass::TensorRef ref_A(block_A.get() + size_t(1) * i * maxM * maxK, Gemm::LayoutA(maxK)); cutlass::TensorRef ref_B(block_B.get() + size_t(1) * i * maxN * maxK, Gemm::LayoutB(maxK)); @@ -320,11 +318,9 @@ struct ExampleRunner { /// Initialize operands to be used in the GEMM and reference GEMM void initialize(ProblemShape const& problem_size) { - auto problem_shape_MNKL = cute::append<4>(problem_size.max_problem_shape, 1); + auto problem_shape_MNKL = cute::append<4>(problem_size.get_host_problem_shape(0), problem_size.groups()); auto [M, N, K, L] = problem_shape_MNKL; - stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, L)); - stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, L)); stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, L)); stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, L)); @@ -340,12 +336,12 @@ struct ExampleRunner { } /// Load a benchmark - std::vector benchmark_problems(std::string const& benchmark_path) { - std::vector problem_sizes_host; + void benchmark_problems(std::string const& benchmark_path) { std::ifstream file(benchmark_path); if (!file.good()) { - return {}; + std::cout << "Issues with benchmark file." << std::endl; + return; } while (file.good()) { @@ -370,30 +366,64 @@ struct ExampleRunner { problem_sizes_host.push_back({extent.m(), extent.n(), extent.k()}); } - return problem_sizes_host; + return; + } + + void benchmark_tokens_per_expert(std::string const& benchmark_path) { + + std::ifstream file(benchmark_path); + if (!file.good()) { + return; + } + + while (file.good()) { + + int idx = -1; + std::string extent_str; + + file >> idx >> extent_str; + + if (idx < 0 || extent_str.empty()) { + break; + } + + cutlass::gemm::GemmCoord extent; + std::vector tokens; + + cutlass::CommandLine::tokenize(tokens, extent_str, 'x'); + + for (int i = 0; i < int(tokens.size()); ++i) { + extent.at(i) = std::atoi(tokens.at(i).c_str()); + } + tokens_per_expert_host.push_back(extent.n()); + } + + return; } bool run(Options const& options, cutlass::KernelHardwareInfo const& hw_info) { - auto problem_sizes_host = benchmark_problems(options.benchmark_path); + + benchmark_problems(options.benchmark_path); if (problem_sizes_host.empty()) { return false; } - problem_sizes.reset(problem_sizes_host.size()); - problem_sizes.copy_from_host(problem_sizes_host.data()); + benchmark_tokens_per_expert(options.benchmark_path); + if (tokens_per_expert_host.empty()) { + return false; + } - ProblemShape problem_size; - problem_size.max_problem_shape = ProblemShapeMax{options.m, options.n, options.k, options.l}; - problem_size.problem_shape.num_groups = problem_sizes_host.size(); - problem_size.problem_shape.problem_shapes = problem_sizes.get(); - problem_size.problem_shape.host_problem_shapes = problem_sizes_host.data(); + tokens_per_expert.reset(tokens_per_expert_host.size()); + tokens_per_expert.copy_from_host(tokens_per_expert_host.data()); + + ProblemShape problem_size {options.m, options.n, options.k, options.groups, tokens_per_expert.get()}; initialize(problem_size); typename Gemm::Arguments arguments{ cutlass::gemm::GemmUniversalMode::kGrouped, problem_size, - {block_A.get(), stride_A, block_B.get(), stride_B}, + {block_A.get(), block_B.get()}, {{}, // epilogue.thread block_C.get(), stride_C, block_D.get(), stride_D}, hw_info @@ -464,7 +494,7 @@ struct ExampleRunner { // Compute average setup and runtime and FLOPs. float elapsed_ms = timer.elapsed_millis(); double avg_runtime_ms = double(elapsed_ms) / double(options.iterations); - double flops = double(int64_t(2) * options.m * options.n * options.k * options.l) / (avg_runtime_ms / 1000.0); + double flops = double(int64_t(2) * options.m * options.n * options.k * options.groups) / (avg_runtime_ms / 1000.0); std::cout << " Avg runtime : " << avg_runtime_ms << " ms" << std::endl; std::cout << " TFLOPS : " << flops / 1e12 << std::endl; diff --git a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_rcgrouped.cu b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_rcgrouped.cu index a6e757d6..6d9f3907 100644 --- a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_rcgrouped.cu +++ b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_rcgrouped.cu @@ -41,7 +41,7 @@ To run this example: - $ ./examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_rcgrouped --m=2048 --n=2048 --k=2048 --groups=10 + $ ./examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_rcgrouped --m=128 --k=128 --groups=10 The above example command makes all 10 groups to be sized at the given m, n, k sizes. Skipping any of the problem dimensions randomizes it across the different groups. @@ -93,7 +93,8 @@ #include "helper.h" using namespace cute; -using ProblemShape = cutlass::gemm::GroupProblemShape>; // per group +using ProblemShape = cutlass::gemm::MoEProblemShape>; // per group + using ElementA = cutlass::float_e4m3_t; // Element type for A matrix operand using ElementB = cutlass::float_e4m3_t; // Element type for B matrix operand using ElementC = cutlass::half_t; // Element type for C and D matrix operands @@ -192,14 +193,11 @@ using StrideB = typename Gemm::GemmKernel::InternalStrideB; using StrideC = typename Gemm::GemmKernel::InternalStrideC; using StrideD = typename Gemm::GemmKernel::InternalStrideD; -StrideA stride_A; - // Host-side allocations std::vector offset_B; std::vector offset_C; std::vector offset_D; -std::vector stride_B_host; std::vector stride_C_host; std::vector stride_D_host; @@ -207,7 +205,7 @@ std::vector alpha_host; std::vector beta_host; // Device-side allocations -cutlass::DeviceAllocation problem_sizes; +cutlass::DeviceAllocation tokens_per_expert; cutlass::DeviceAllocation block_A; cutlass::DeviceAllocation block_B; @@ -223,7 +221,6 @@ cutlass::DeviceAllocation ptr_C; cutlass::DeviceAllocation ptr_D; cutlass::DeviceAllocation ptr_ref_D; -cutlass::DeviceAllocation stride_B; cutlass::DeviceAllocation stride_C; cutlass::DeviceAllocation stride_D; @@ -251,13 +248,14 @@ struct Options { float beta = FLT_MAX; int iterations = 1000; int warmup = 1000; - int m = 1024, n = 2048, k = 512, groups = 10; + int m = 128, n = 128, k = 128, groups = 10; double sparse_prob = 0.1; dim3 cluster_shape = dim3(4,2,1); dim3 cluster_shape_fallback = dim3(2,1,1); RasterOrderOptions raster_order = RasterOrderOptions::AlongM; int max_sm_count = INT_MAX; std::string benchmark_path; + std::vector tokens_per_expert_host; std::vector problem_sizes_host; int const tma_alignment_bits = 128; int const alignment = tma_alignment_bits / cutlass::sizeof_bits::value; @@ -342,6 +340,7 @@ struct Options { n = alignment * ((rand() % 64) + 1); } problem_sizes_host.push_back({m, n, k}); + tokens_per_expert_host.push_back(n); } } @@ -512,8 +511,6 @@ bool initialize_block( scope_min = static_cast(-8); } - scope_min = static_cast(0); - scope_max = static_cast(2); cutlass::reference::device::BlockFillRandomUniform( block.get(), block.size(), seed, scope_max, scope_min, 0); @@ -545,7 +542,6 @@ void allocate(const Options &options) { total_elements_C += elements_C; total_elements_D += elements_D; - stride_B_host.push_back(cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1})); stride_C_host.push_back(cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1})); stride_D_host.push_back(cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1})); @@ -558,7 +554,6 @@ void allocate(const Options &options) { block_alpha.reset(options.groups); block_beta.reset(options.groups); - stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(options.m, options.k, options.groups)); auto a_coord = cutlass::make_Coord(options.m * options.groups, options.k); block_A.reset(a_coord.product()); @@ -569,8 +564,9 @@ void initialize(const Options &options) { uint64_t seed = 2020; - problem_sizes.reset(options.groups); - problem_sizes.copy_from_host(options.problem_sizes_host.data()); + // Setting up tokens_per_expert array + tokens_per_expert.reset(options.tokens_per_expert_host.size()); + tokens_per_expert.copy_from_host(options.tokens_per_expert_host.data()); // // Assign pointers @@ -601,9 +597,6 @@ void initialize(const Options &options) { ptr_D.reset(options.groups); ptr_D.copy_from_host(ptr_D_host.data()); - stride_B.reset(options.groups); - stride_B.copy_from_host(stride_B_host.data()); - stride_C.reset(options.groups); stride_C.copy_from_host(stride_C_host.data()); @@ -675,12 +668,12 @@ typename Gemm::Arguments args_from_options(Options &options) typename Gemm::GemmKernel::TileSchedulerArguments scheduler; scheduler.raster_order = options.raster_order; - + arguments = typename Gemm::Arguments { cutlass::gemm::GemmUniversalMode::kGrouped, - {options.groups, problem_sizes.get(), options.problem_sizes_host.data()}, - {block_A.get(), stride_A, ptr_B.get(), stride_B.get()}, - {fusion_args, ptr_C.get(), stride_C.get(), ptr_D.get(), stride_D.get()}, + {options.m, options.n, options.k, options.groups, tokens_per_expert.get()}, + {block_A.get(), ptr_B.get()}, + {fusion_args, ptr_C.get(), nullptr, ptr_D.get(), nullptr}, hw_info, scheduler }; diff --git a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_regular.cu b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_regular.cu index f33838c8..754312e0 100644 --- a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_regular.cu +++ b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_regular.cu @@ -39,7 +39,7 @@ Usage: $ ./examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_regular - --m=28672 --n=4 --k=4096 --l=8 + --m=7168 --n=128 --k=512 --l=8 */ @@ -320,19 +320,38 @@ struct ExampleRunner { initialize_block(block_C, seed + 2021); } + + /// Populates a Gemm::Arguments structure from the given commandline options + typename Gemm::Arguments args_from_options(ProblemShapeType problem_size, cutlass::KernelHardwareInfo const& hw_info) + { + if constexpr (std::is_same_v){ + return typename Gemm::Arguments { + cutlass::gemm::GemmUniversalMode::kGemm, + problem_size, + {block_A.get(), block_B.get()}, + {{}, // epilogue.thread + block_C.get(), stride_C, block_D.get(), stride_D}, + hw_info + }; + } + else { + return typename Gemm::Arguments { + cutlass::gemm::GemmUniversalMode::kGemm, + problem_size, + {block_A.get(), stride_A, block_B.get(), stride_B}, + {{}, // epilogue.thread + block_C.get(), stride_C, block_D.get(), stride_D}, + hw_info + }; + } + } + bool run(Options const& options, cutlass::KernelHardwareInfo const& hw_info) { ProblemShapeType problem_size = ProblemShapeType{options.m, options.n, options.k, options.l}; initialize(problem_size); - typename Gemm::Arguments arguments{ - cutlass::gemm::GemmUniversalMode::kGemm, - problem_size, - {block_A.get(), stride_A, block_B.get(), stride_B}, - {{}, // epilogue.thread - block_C.get(), stride_C, block_D.get(), stride_D}, - hw_info - }; + typename Gemm::Arguments arguments = args_from_options(problem_size, hw_info); // arguments.scheduler.max_swizzle_size = options.swizzle; @@ -471,7 +490,7 @@ int main(int argc, char const **args) { runner_tma.run(options, hw_info); std::cout << "Running kernel with CPASYNC load:" << std::endl; - ExampleRunner runner_cpasync; + ExampleRunner runner_cpasync; runner_cpasync.run(options, hw_info); std::cout << "Running kernel with mixed TMA+CPASYNC load:" << std::endl; diff --git a/examples/92_blackwell_moe_gemm/CMakeLists.txt b/examples/92_blackwell_moe_gemm/CMakeLists.txt index b0ece023..dbdd5313 100644 --- a/examples/92_blackwell_moe_gemm/CMakeLists.txt +++ b/examples/92_blackwell_moe_gemm/CMakeLists.txt @@ -26,18 +26,17 @@ # 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(TEST_MIXTRAL_A --m=28672 --n=4 --k=4096 --l=8) -set(TEST_MIXTRAL_B --m=4096 --n=4 --k=14336 --l=8) -set(TEST_DEEPSEEK_A --m=4096 --n=1 --k=7168 --l=256) -set(TEST_DEEPSEEK_B --m=7168 --n=1 --k=2048 --l=256) -set(TEST_IRREGULAR_MNK --m=4080 --n=9 --k=4112 --l=8) # M,N,K not multiples of tile size +set(TEST_MIXTRAL_A --m=28672 --n=4 --k=4096 --groups=8 --iterations=0) +set(TEST_MIXTRAL_B --m=4096 --n=4 --k=14336 --groups=8 --iterations=0) +set(TEST_DEEPSEEK_A --m=4096 --n=1 --k=7168 --groups=256 --iterations=0) +set(TEST_DEEPSEEK_B --m=7168 --n=1 --k=2048 --groups=256 --iterations=0) +set(TEST_IRREGULAR_MNK --m=4080 --n=9 --k=4112 --groups=8 --iterations=0) # M,N,K not mugroupstipgroupses of tigroupse size -set(TEST_DEEPSEEK_A_FP4 --m=1024 --n=1 --k=7168 --l=256) # TP=1 shape is too large for PackedVectorLayout -set(TEST_DEEPSEEK_B_FP4 --m=7168 --n=1 --k=512 --l=256) -set(TEST_IRREGULAR_MNK_FP4 --m=4080 --n=9 --k=4160 --l=8) +set(TEST_DEEPSEEK_A_FP4 --m=1024 --n=1 --k=7168 --groups=256 --iterations=0) # TP=1 shape is too groupsarge for PackedVectorLayout +set(TEST_DEEPSEEK_B_FP4 --m=7168 --n=1 --k=512 --groups=256 --iterations=0) +set(TEST_IRREGULAR_MNK_FP4 --m=4080 --n=9 --k=4160 --groups=8 --iterations=0) set(TEST_FIXED --m=2048 --n=5120 --k=8192 --iterations=0) # Fixed problem sizes -set(TEST_SPARSE_GROUPS --sparse_test --sparse_prob=0.3 --iterations=0) if (CUTLASS_NVCC_ARCHS MATCHES 100a) cutlass_example_add_executable( @@ -61,7 +60,6 @@ cutlass_example_add_executable( 92_blackwell_moe_gemm_rcgrouped.cu TEST_COMMAND_OPTIONS TEST_FIXED - TEST_SPARSE_GROUPS ) cutlass_example_add_executable( diff --git a/examples/python/CuTeDSL/ampere/call_bypass_dlpack.py b/examples/python/CuTeDSL/ampere/call_bypass_dlpack.py index 6f2504f7..fbae6bfd 100644 --- a/examples/python/CuTeDSL/ampere/call_bypass_dlpack.py +++ b/examples/python/CuTeDSL/ampere/call_bypass_dlpack.py @@ -74,9 +74,11 @@ with stride-1 which propagate alignment incorrectly. """ -# Add the current directory to sys.path -sys.path.append(os.path.dirname(os.path.abspath(__file__))) -from tensorop_gemm import TensorOpGemm +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "..")) + +from ampere.tensorop_gemm import TensorOpGemm @cute.jit diff --git a/examples/python/CuTeDSL/ampere/call_from_jit.py b/examples/python/CuTeDSL/ampere/call_from_jit.py index 68e1b0ae..30c759ce 100644 --- a/examples/python/CuTeDSL/ampere/call_from_jit.py +++ b/examples/python/CuTeDSL/ampere/call_from_jit.py @@ -67,10 +67,11 @@ import cutlass.cute as cute from cutlass.torch import dtype as torch_dtype from cutlass.cute.runtime import make_ptr +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "..")) -# Add the current directory to sys.path -sys.path.append(os.path.dirname(os.path.abspath(__file__))) -from tensorop_gemm import TensorOpGemm +from ampere.tensorop_gemm import TensorOpGemm class BufferWithLayout: diff --git a/examples/python/CuTeDSL/ampere/call_with_tvm_ffi.py b/examples/python/CuTeDSL/ampere/call_with_tvm_ffi.py new file mode 100644 index 00000000..39319a34 --- /dev/null +++ b/examples/python/CuTeDSL/ampere/call_with_tvm_ffi.py @@ -0,0 +1,128 @@ +# 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 sys +import os +import torch +import time + +import cutlass +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack + +"""Demonstrates calling off-the-shelf kernels with TVM FFI without DLPack. + +This example shows how to compile CuTe JIT function with fake tensors then run it with TVM FFI. +""" + +if __name__ == "__main__": + # Add the current directory to sys.path + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "..")) +from ampere.tensorop_gemm import TensorOpGemm + + +def compile_op(use_tvm_ffi: bool = True): + from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_tensor + + a_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int()) + b_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int()) + c_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int()) + a = make_fake_compact_tensor( + cutlass.Float16, a_shape, stride_order=(1, 0, 2), assumed_align=16 + ) + b = make_fake_compact_tensor( + cutlass.Float16, b_shape, stride_order=(1, 0, 2), assumed_align=16 + ) + c = make_fake_compact_tensor( + cutlass.Float16, c_shape, stride_order=(1, 0, 2), assumed_align=16 + ) + + tensor_op_gemm = TensorOpGemm( + cutlass.Float16, cutlass.Float16, cutlass.Float32, (2, 2, 1) + ) + compiled_fn = cute.compile( + tensor_op_gemm, a, b, c, options="--enable-tvm-ffi" if use_tvm_ffi else "" + ) + return compiled_fn + + +def run_op(compiled_fn, mnkl, *, use_tvm_ffi: bool = True): + print("\nRunning TensorOpGemm test with:") + print(f"Tensor dimensions: {mnkl}") + torch.manual_seed(1112) + # (M,K,L) + a_torch = torch.randn( + mnkl[3], mnkl[0], mnkl[2], dtype=torch.float16, device="cuda" + ).permute(1, 2, 0) + # (N,K,L) + b_torch = torch.randn( + mnkl[3], mnkl[1], mnkl[2], dtype=torch.float16, device="cuda" + ).permute(1, 2, 0) + # (N,M,L) + c_torch = torch.randn( + mnkl[3], mnkl[0], mnkl[1], dtype=torch.float16, device="cuda" + ).permute(1, 2, 0) + + print("Input tensor shapes:") + print(f"a: {a_torch.shape}, dtype: {a_torch.dtype}") + print(f"b: {b_torch.shape}, dtype: {b_torch.dtype}") + print(f"c: {c_torch.shape}, dtype: {c_torch.dtype}\n") + if not use_tvm_ffi: + a = from_dlpack(a_torch).mark_layout_dynamic(leading_dim=1) + b = from_dlpack(b_torch).mark_layout_dynamic(leading_dim=1) + c = from_dlpack(c_torch).mark_layout_dynamic(leading_dim=1) + else: + a = a_torch + b = b_torch + c = c_torch + # pass in torch tensor as input + compiled_fn(a, b, c) + torch.cuda.synchronize() + + # measure the launch overhead of tvm ffi function + repeat = 100 + start_time = time.time() + for i in range(repeat): + compiled_fn(a, b, c) + end_time = time.time() + print( + f"Launch overhead of tvm ffi function: {(end_time - start_time) / repeat} seconds" + ) + + ref = torch.einsum("mkl,nkl->mnl", a_torch, b_torch) + torch.testing.assert_close(c_torch, ref, atol=1e-05, rtol=1e-05) + print("\n[DSL INFO] Results verified successfully!") + print(f"First few elements of result: \n{c_torch[:3, :3, :3]}") + + +if __name__ == "__main__": + compiled_fn = compile_op(use_tvm_ffi=False) + run_op(compiled_fn, [512, 512, 256, 1], use_tvm_ffi=False) + compiled_fn = compile_op(use_tvm_ffi=True) + run_op(compiled_fn, [512, 512, 256, 1], use_tvm_ffi=True) diff --git a/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py b/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py index 7d76f724..a5be849d 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py @@ -39,6 +39,7 @@ 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 import cutlass.utils.blackwell_helpers as sm100_utils import math @@ -220,22 +221,18 @@ class BlockwiseGemmKernel: self.num_regs_epilogue_warps = 216 self.num_regs_acc_update_warps = 216 - # Set barrier for cta sync, epilogue sync and tmem ptr sync - self.cta_sync_barrier = pipeline.NamedBarrier( - barrier_id=1, - num_threads=self.threads_per_cta, - ) + # Set barrier for epilogue sync and tmem ptr sync self.epilog_sync_barrier = pipeline.NamedBarrier( - barrier_id=2, + barrier_id=1, num_threads=32 * len(self.epilog_warp_id), ) self.tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=3, + barrier_id=2, num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id, *self.acc_update_warp_id)), ) self.sched_sync_barrier = pipeline.NamedBarrier( - barrier_id=4, + barrier_id=3, num_threads=self.threads_per_warp, ) self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") @@ -703,6 +700,7 @@ class BlockwiseGemmKernel: consumer_group=ab_pipeline_consumer_group, tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, ) # Initialize mainloop scale_pipeline (barrier) and states @@ -719,6 +717,7 @@ class BlockwiseGemmKernel: num_stages=self.num_scale_stage, producer_group=scale_pipeline_producer_group, consumer_group=scale_pipeline_consumer_group, + defer_sync=True, ) # Initialize acc_pipeline (barrier) and states @@ -735,6 +734,7 @@ class BlockwiseGemmKernel: producer_group=acc_pipeline_producer_group, consumer_group=acc_pipeline_consumer_group, cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, ) # Initialize epilogue pipeline (barrier) and states @@ -751,6 +751,7 @@ class BlockwiseGemmKernel: num_stages=1, producer_group=epi_pipeline_producer_group, consumer_group=epi_pipeline_consumer_group, + defer_sync=True, ) # Initialize tile info pipeline (barrier) and states @@ -767,6 +768,7 @@ class BlockwiseGemmKernel: num_stages=self.num_tile_stage, producer_group=tile_info_pipeline_producer_group, consumer_group=tile_info_pipeline_consumer_group, + defer_sync=True, ) # Tensor memory dealloc barrier init @@ -779,8 +781,7 @@ class BlockwiseGemmKernel: ) # Cluster arrive after barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # # Setup smem tensor A/B/C/Scale @@ -968,10 +969,7 @@ class BlockwiseGemmKernel: # # Cluster wait before tensor memory alloc # - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_wait() - else: - self.cta_sync_barrier.arrive_and_wait() + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) # # Specialized Schedule warp 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 47d89f0d..43b29d06 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py @@ -39,6 +39,7 @@ 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 import cutlass.utils.blackwell_helpers as sm100_utils from cutlass.cute.runtime import from_dlpack @@ -236,22 +237,18 @@ class BlockwiseContiguousGroupedGemmKernel: self.num_regs_epilogue_warps = 216 self.num_regs_acc_update_warps = 216 - # Set barrier for cta sync, epilogue sync and tmem ptr sync - self.cta_sync_barrier = pipeline.NamedBarrier( - barrier_id=1, - num_threads=self.threads_per_cta, - ) + # Set barrier for epilogue sync and tmem ptr sync self.epilog_sync_barrier = pipeline.NamedBarrier( - barrier_id=2, + barrier_id=1, num_threads=32 * len(self.epilog_warp_id), ) self.tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=3, + barrier_id=2, num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id, *self.acc_update_warp_id)), ) self.sched_sync_barrier = pipeline.NamedBarrier( - barrier_id=4, + barrier_id=3, num_threads=self.threads_per_warp, ) self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") @@ -724,6 +721,7 @@ class BlockwiseContiguousGroupedGemmKernel: consumer_group=ab_pipeline_consumer_group, tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, ) # Initialize mainloop scale_pipeline (barrier) and states @@ -740,6 +738,7 @@ class BlockwiseContiguousGroupedGemmKernel: num_stages=self.num_scale_stage, producer_group=scale_pipeline_producer_group, consumer_group=scale_pipeline_consumer_group, + defer_sync=True, ) # Initialize acc_pipeline (barrier) and states @@ -756,6 +755,7 @@ class BlockwiseContiguousGroupedGemmKernel: producer_group=acc_pipeline_producer_group, consumer_group=acc_pipeline_consumer_group, cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, ) # Initialize epilogue pipeline (barrier) and states @@ -772,6 +772,7 @@ class BlockwiseContiguousGroupedGemmKernel: num_stages=1, producer_group=epi_pipeline_producer_group, consumer_group=epi_pipeline_consumer_group, + defer_sync=True, ) # Initialize tile info pipeline (barrier) and states @@ -788,6 +789,7 @@ class BlockwiseContiguousGroupedGemmKernel: num_stages=self.num_tile_stage, producer_group=tile_info_pipeline_producer_group, consumer_group=tile_info_pipeline_consumer_group, + defer_sync=True, ) # Tensor memory dealloc barrier init @@ -800,8 +802,7 @@ class BlockwiseContiguousGroupedGemmKernel: ) # Cluster arrive after barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # # Setup smem tensor A/B/C/Scale @@ -989,10 +990,7 @@ class BlockwiseContiguousGroupedGemmKernel: # # Cluster wait before tensor memory alloc # - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_wait() - else: - self.cta_sync_barrier.arrive_and_wait() + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) # # Specialized Schedule warp 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 a08285d9..2457ea04 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py @@ -39,6 +39,7 @@ 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 import cutlass.utils.blackwell_helpers as sm100_utils from cutlass.cute.runtime import from_dlpack @@ -235,22 +236,18 @@ class BlockwiseMaskedGroupedGemmKernel: self.num_regs_epilogue_warps = 216 self.num_regs_acc_update_warps = 216 - # 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, - ) + # Set barrier id for epilogue sync and tmem ptr sync self.epilog_sync_barrier = pipeline.NamedBarrier( - barrier_id=2, + barrier_id=1, num_threads=32 * len(self.epilog_warp_id), ) self.tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=3, + barrier_id=2, num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id, *self.acc_update_warp_id)), ) self.sched_sync_barrier = pipeline.NamedBarrier( - barrier_id=4, + barrier_id=3, num_threads=self.threads_per_warp, ) self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") @@ -723,6 +720,7 @@ class BlockwiseMaskedGroupedGemmKernel: consumer_group=ab_pipeline_consumer_group, tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, ) # Initialize mainloop scale_pipeline (barrier) and states @@ -739,6 +737,7 @@ class BlockwiseMaskedGroupedGemmKernel: num_stages=self.num_scale_stage, producer_group=scale_pipeline_producer_group, consumer_group=scale_pipeline_consumer_group, + defer_sync=True, ) # Initialize acc_pipeline (barrier) and states @@ -755,6 +754,7 @@ class BlockwiseMaskedGroupedGemmKernel: producer_group=acc_pipeline_producer_group, consumer_group=acc_pipeline_consumer_group, cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, ) # Initialize epilogue pipeline (barrier) and states @@ -771,6 +771,7 @@ class BlockwiseMaskedGroupedGemmKernel: num_stages=1, producer_group=epi_pipeline_producer_group, consumer_group=epi_pipeline_consumer_group, + defer_sync=True, ) # Initialize tile info pipeline (barrier) and states @@ -787,6 +788,7 @@ class BlockwiseMaskedGroupedGemmKernel: num_stages=self.num_tile_stage, producer_group=tile_info_pipeline_producer_group, consumer_group=tile_info_pipeline_consumer_group, + defer_sync=True, ) # Tensor memory dealloc barrier init @@ -799,8 +801,7 @@ class BlockwiseMaskedGroupedGemmKernel: ) # Cluster arrive after barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # # Setup smem tensor A/B/C/Scale @@ -988,10 +989,7 @@ class BlockwiseMaskedGroupedGemmKernel: # # Cluster wait before tensor memory alloc # - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_wait() - else: - self.cta_sync_barrier.arrive_and_wait() + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) # # Specialized Schedule warp diff --git a/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py b/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py index 84367506..c37d094f 100644 --- a/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py +++ b/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py @@ -38,6 +38,7 @@ 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 import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.utils.blockscaled_layout as blockscaled_utils from cutlass.cute.runtime import from_dlpack @@ -208,17 +209,13 @@ class Sm100BlockScaledPersistentDenseGemmKernel: 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, - ) + # Set barrier id for epilogue sync and tmem ptr sync self.epilog_sync_barrier = pipeline.NamedBarrier( - barrier_id=2, + barrier_id=1, num_threads=32 * len(self.epilog_warp_id), ) self.tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=3, + barrier_id=2, num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)), ) self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") @@ -288,6 +285,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel: self.mma_tiler[1], self.mma_tiler[2], ) + self.cta_tile_shape_mnk_sfb = ( + self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_sfb[1], + self.mma_tiler_sfb[2], + ) # Compute cluster layout self.cluster_layout_vmnk = cute.tiled_divide( @@ -314,6 +316,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: self.c_layout, self.c_dtype, ) + self.epi_tile_n = cute.size(self.epi_tile[1]) # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( @@ -362,6 +365,19 @@ class Sm100BlockScaledPersistentDenseGemmKernel: 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 + + # 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_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 + + # Only when overlapping_accum is enabled, we need to release accumulator buffer early in epilogue + self.iter_acc_early_release_in_epilogue = self.num_sf_tmem_cols // self.epi_tile_n + @cute.jit def __call__( self, @@ -640,6 +656,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: block=[self.threads_per_cta, 1, 1], cluster=(*self.cluster_shape_mn, 1), stream=stream, + min_blocks_per_mp=1, ) return @@ -726,6 +743,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: 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 @@ -742,6 +760,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: 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 @@ -754,8 +773,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: ) # Cluster arrive after barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # # Setup smem tensor A/B/SFA/SFB/C @@ -910,18 +928,34 @@ class Sm100BlockScaledPersistentDenseGemmKernel: 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) - ) + 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], + (256 - 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 # - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_wait() - else: - self.cta_sync_barrier.arrive_and_wait() + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) # # Specialized TMA load warp @@ -1057,7 +1091,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: # Make SFA tmem tensor sfa_tmem_ptr = cute.recast_ptr( - acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base), + acc_tmem_ptr + self.num_accumulator_tmem_cols, dtype=self.sf_dtype, ) # (MMA, MMA_M, MMA_K) @@ -1071,9 +1105,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: # 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), + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, dtype=self.sf_dtype, ) # (MMA, MMA_N, MMA_K) @@ -1122,9 +1154,15 @@ class Sm100BlockScaledPersistentDenseGemmKernel: 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 # (MMA, MMA_M, MMA_N) - tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)] # Peek (try_wait) AB buffer full for k_tile = 0 ab_consumer_state.reset_count() @@ -1146,8 +1184,8 @@ class Sm100BlockScaledPersistentDenseGemmKernel: offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0) shifted_ptr = cute.recast_ptr( acc_tmem_ptr - + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base) - + tcgen05.find_tmem_tensor_col_offset(tCtSFA) + + self.num_accumulator_tmem_cols + + self.num_sfa_tmem_cols + offset, dtype=self.sf_dtype, ) @@ -1156,9 +1194,9 @@ 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 - + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base) - + tcgen05.find_tmem_tensor_col_offset(tCtSFA) + acc_tmem_ptr + + self.num_accumulator_tmem_cols + + self.num_sfa_tmem_cols + offset, dtype=self.sf_dtype, ) @@ -1350,10 +1388,17 @@ 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) + 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_consumer_state.index) + (None, None, None, None, None, acc_stage_index) ] # @@ -1370,12 +1415,27 @@ class Sm100BlockScaledPersistentDenseGemmKernel: 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): + 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 # # Load accumulator from tensor memory buffer to register # - tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)] cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + # + # Async arrive accumulator buffer empty ealier when overlapping_accum is enabled + # + if cutlass.const_expr(self.overlapping_accum): + if subtile_idx == self.iter_acc_early_release_in_epilogue: + # Fence for TMEM load + cute.arch.fence_view_async_tmem_load() + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + # # Convert to C type # @@ -1386,7 +1446,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: # # Store C to shared memory # - c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage + c_buffer = (num_prev_subtiles + real_subtile_idx) % self.num_c_stage cute.copy( tiled_copy_r2s, tRS_rC, @@ -1406,7 +1466,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: cute.copy( tma_atom_c, bSG_sC[(None, c_buffer)], - bSG_gC[(None, subtile_idx)], + bSG_gC[(None, real_subtile_idx)], ) # Fence and barrier to make sure shared memory store is visible to TMA store c_pipeline.producer_commit() @@ -1416,9 +1476,10 @@ class Sm100BlockScaledPersistentDenseGemmKernel: # # Async arrive accumulator buffer empty # - with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() + if cutlass.const_expr(not self.overlapping_accum): + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() # # Advance to next tile @@ -2286,6 +2347,7 @@ def run( c_tensor, max_active_clusters, current_stream, + options=f"--opt-level 2", ) # Compute reference result diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm.py b/examples/python/CuTeDSL/blackwell/dense_gemm.py index 304c86bb..4f2e93ea 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm.py @@ -36,6 +36,7 @@ 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 @@ -536,6 +537,7 @@ class DenseGemmKernel: 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 @@ -549,6 +551,7 @@ class DenseGemmKernel: producer_group=acc_pipeline_producer_group, consumer_group=acc_pipeline_consumer_group, cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, ) acc_producer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Producer, self.num_acc_stage @@ -569,8 +572,7 @@ class DenseGemmKernel: ) # Cluster arrive after barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) # # Setup smem tensor A/B/C @@ -686,8 +688,7 @@ class DenseGemmKernel: # # Cluster wait before tensor memory alloc # - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_wait() + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) # Alloc tensor memory buffer tmem.allocate(self.num_tmem_alloc_cols) 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 c6262f91..712d9eb0 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py @@ -38,6 +38,7 @@ 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 import cutlass.utils.blackwell_helpers as sm100_utils from cutlass.cute.nvgpu import cpasync, tcgen05 @@ -232,9 +233,8 @@ class SM100PersistentDenseGemmAlphaBetaKernel: ) ) # Set barrier id for cta sync, epilogue sync and tmem ptr sync - self.cta_sync_bar_id = 1 - self.epilog_sync_bar_id = 2 - self.tmem_alloc_sync_bar_id = 3 + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") def _setup_attributes(self): @@ -635,6 +635,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel: 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 @@ -651,6 +652,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel: producer_group=acc_pipeline_producer_group, consumer_group=acc_pipeline_consumer_group, cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, ) # Load C pipeline @@ -665,6 +667,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel: producer_group=c_producer_group, consumer_group=c_consumer_group, tx_count=self.tma_c_load_bytes, + defer_sync=True, ) tmem_alloc_barrier = pipeline.NamedBarrier( @@ -681,8 +684,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel: ) # Cluster arrive after barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # # Setup smem tensor A/B/C/D @@ -796,9 +798,6 @@ class SM100PersistentDenseGemmAlphaBetaKernel: # Named barriers # - cta_sync_barrier = pipeline.NamedBarrier( - self.cta_sync_bar_id, self.threads_per_cta - ) epilog_sync_barrier = pipeline.NamedBarrier( self.epilog_sync_bar_id, 32 * len(self.epilog_warp_ids) ) @@ -806,10 +805,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel: # # 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() + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) # # Specialized TMA load warp diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py index 2ff8ff26..d5b48f45 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py @@ -29,16 +29,14 @@ 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 -import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.cute.nvgpu import cpasync, tcgen05 """ @@ -152,10 +150,10 @@ def _compute_stages( num_c_stage = 2 if use_tma_store else 0 # Calculate smem layout and size for one stage of A, B, and C with 1-stage - a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + a_smem_layout_stage_one = utils.sm100.make_smem_layout_a( tiled_mma, mma_tiler_mnk, a_dtype, 1 ) - b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + b_smem_layout_staged_one = utils.sm100.make_smem_layout_b( tiled_mma, mma_tiler_mnk, b_dtype, 1 ) @@ -229,14 +227,14 @@ class PersistentDenseGemmKernel: - Cluster shape M must be multiple of 2 if use_2cta_instrs=True - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 - Example: - >>> gemm = PersistentDenseGemmKernel( - ... acc_dtype=cutlass.Float32, - ... use_2cta_instrs=True, - ... mma_tiler_mn=(128, 128), - ... cluster_shape_mn=(2, 2) - ... ) - >>> gemm(a_tensor, b_tensor, c_tensor, max_active_clusters, stream) + **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__( @@ -316,7 +314,7 @@ class PersistentDenseGemmKernel: - Computing tensor memory allocation columns """ # Configure tiled mma - tiled_mma = sm100_utils.make_trivial_tiled_mma( + tiled_mma = utils.sm100.make_trivial_tiled_mma( self.a_dtype, self.a_major_mode, self.b_major_mode, @@ -353,7 +351,7 @@ class PersistentDenseGemmKernel: # Compute epilogue subtile if cutlass.const_expr(self.use_tma_store): - self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.epi_tile = utils.sm100.compute_epilogue_tile_shape( self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, @@ -364,7 +362,7 @@ class PersistentDenseGemmKernel: c_smem_layout = None if cutlass.const_expr(self.use_tma_store): - c_smem_layout = sm100_utils.make_smem_layout_epi( + c_smem_layout = utils.sm100.make_smem_layout_epi( self.c_dtype, self.c_layout, self.epi_tile, 1 ) @@ -382,16 +380,16 @@ class PersistentDenseGemmKernel: ) # Compute A/B/C shared memory layout - self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + 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 = sm100_utils.make_smem_layout_b( + 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 = sm100_utils.make_smem_layout_epi( + self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage ) @@ -447,7 +445,7 @@ class PersistentDenseGemmKernel: # Setup attributes that dependent on gemm inputs self._setup_attributes() - tiled_mma = sm100_utils.make_trivial_tiled_mma( + tiled_mma = utils.sm100.make_trivial_tiled_mma( self.a_dtype, self.a_major_mode, self.b_major_mode, @@ -458,7 +456,7 @@ class PersistentDenseGemmKernel: 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( + a_op = utils.sm100.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)) @@ -475,7 +473,7 @@ class PersistentDenseGemmKernel: ) # Setup TMA load for B - b_op = sm100_utils.cluster_shape_to_tma_atom_B( + b_op = utils.sm100.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)) @@ -614,6 +612,7 @@ class PersistentDenseGemmKernel: 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 @@ -630,6 +629,7 @@ class PersistentDenseGemmKernel: 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( @@ -652,8 +652,7 @@ class PersistentDenseGemmKernel: ) # Cluster arrive after barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) # # Setup smem tensor A/B/C @@ -761,10 +760,7 @@ class PersistentDenseGemmKernel: # # Cluster wait before tensor memory alloc # - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_wait() - else: - cute.arch.sync_threads() + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) # # Specialized TMA load warp @@ -1297,7 +1293,7 @@ class PersistentDenseGemmKernel: :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] """ # Make tiledCopy for tensor memory load - copy_atom_t2r = sm100_utils.get_tmem_load_op( + copy_atom_t2r = utils.sm100.get_tmem_load_op( self.cta_tile_shape_mnk, self.c_layout, self.c_dtype, @@ -1354,7 +1350,7 @@ class PersistentDenseGemmKernel: - 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( + 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) @@ -1617,97 +1613,187 @@ class PersistentDenseGemmKernel: is_valid = False return is_valid - def can_implement(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor) -> bool: - """Check if the given tensors can be implemented by this kernel. + def can_implement( + self, + mnkl: Tuple[int, int, int, int], + ab_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 a: Input tensor A - :type a: cute.Tensor - :param b: Input tensor B - :type b: cute.Tensor - :param c: Output tensor C - :type c: cute.Tensor - - :return: True if the gemm supports the given config, False otherwise + :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 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 """ - m, n, k, l = a.shape[0], b.shape[0], a.shape[1], a.shape[2] - # infer a_major, b_major, c_major - is_m_major_a = utils.LayoutEnum.from_tensor(a).is_m_major_a() - is_n_major_b = utils.LayoutEnum.from_tensor(b).is_n_major_b() - is_m_major_c = utils.LayoutEnum.from_tensor(c).is_m_major_c() - a_major = "m" if is_m_major_a else "k" - b_major = "n" if is_n_major_b else "k" - c_major = "m" if is_m_major_c else "n" - - can_implement = True # Skip unsupported types - if not self.is_valid_dtypes(a.element_type, c.element_type): - can_implement = False + if not self.is_valid_dtypes(ab_dtype, c_dtype): + return False + # Skip invalid mma tile shape and cluster shape if not self.is_valid_mma_tiler_and_cluster_shape(): - can_implement = False + return False + + # 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, a.element_type, c.element_type, a_major, b_major, c_major + m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major ): - can_implement = False + return False # Skip invalid epilogue store option if not self.is_valid_epilog_store_option(m, n): - can_implement = False + return False - return can_implement + return True -def create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype): - torch.manual_seed(1111) +@cute.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). - 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) - c_torch_cpu = cutlass_torch.matrix(l, m, n, c_major == "m", c_dtype) + Internally, the tensors are permuted to match CuTe's convention: + - a: (m, k, l) + - b: (n, k, l) + - c: (m, n, l) - a_tensor, _ = cutlass_torch.cute_tensor_like( - a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + :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) + + +def compile_bmm( + gemm_op: PersistentDenseGemmKernel, + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + options: str = "", +): + from cutlass.cute.runtime import make_fake_compact_tensor + + a_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int()) + b_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int()) + c_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int()) + + if a_major == "k": + a_order = (2, 1, 0) # k is leading dimension + elif a_major == "m": + a_order = (2, 0, 1) # m is leading dimension + + if b_major == "n": + b_order = (2, 1, 0) # n is leading dimension + elif b_major == "k": + b_order = (2, 0, 1) # k is leading dimension + + if c_major == "n": + c_order = (2, 1, 0) # n is leading dimension + elif c_major == "m": + c_order = (2, 0, 1) # m is leading dimension + + a = make_fake_compact_tensor( + a_dtype, a_shape, stride_order=a_order, assumed_align=16 ) - b_tensor, _ = cutlass_torch.cute_tensor_like( - b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + b = make_fake_compact_tensor( + b_dtype, b_shape, stride_order=b_order, assumed_align=16 ) - c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like( - c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16 + c = make_fake_compact_tensor( + c_dtype, c_shape, stride_order=c_order, assumed_align=16 ) + return cute.compile( + bmm, gemm_op, a, b, c, max_active_clusters, stream, epilogue_op, options=options + ) + + +def prepare_tensors( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + init_random: bool = True, +): + 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") + elif a_major == "m": + a = 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") + elif b_major == "k": + b = 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") + elif c_major == "m": + c = 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) + return ( - a_tensor, - b_tensor, - c_tensor, - a_torch_cpu, - b_torch_cpu, - c_torch_cpu, - c_torch_gpu, + a.to(dtype=torch_dtype(ab_dtype)), + b.to(dtype=torch_dtype(ab_dtype)), + c.to(dtype=torch_dtype(c_dtype)), ) -def compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance): - # Copy gpu result back - kernel_result = c_torch_gpu.cpu() - - # Compute reference result - ref = torch.einsum( - "mkl,nkl->mnl", - a_torch_cpu.to(dtype=torch.float32), - b_torch_cpu.to(dtype=torch.float32), - ) - - # 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( mnkl: Tuple[int, int, int, int], ab_dtype: Type[cutlass.Numeric], @@ -1725,48 +1811,55 @@ def run( iterations: int = 1, skip_ref_check: bool = False, use_cold_l2: bool = False, + use_tvm_ffi: bool = False, + benchmark: bool = False, **kwargs, ): - """Execute a persistent batched dense GEMM operation on Blackwell architecture with performance benchmarking. + """ + Execute a persistent batched dense GEMM operation on Blackwell architecture with performance benchmarking. - This function prepares input tensors, configures and launches the persistent GEMM kernel, - optionally performs reference validation, and benchmarks the execution performance. + Prepares input tensors, configures and launches the persistent GEMM kernel, + optionally performs reference validation, and benchmarks execution. - :param mnkl: Problem size (M, N, K, L) + :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 + :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 + :param c_dtype: Data type for output tensor C. :type c_dtype: Type[cutlass.Numeric] - :param acc_dtype: Data type for accumulation during matrix multiplication + :param acc_dtype: Accumulator data type for the matrix multiplication. :type acc_dtype: Type[cutlass.Numeric] - :param a_major/b_major/c_major: Memory layout of tensor A/B/C - :type a_major/b_major/c_major: str - :param mma_tiler_mn: MMA tiling size. If not specified in the decorator parameters, the autotuner will use the - default value of (256, 256). Otherwise, the autotuner will use the value specified in the decorator parameters. + :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. If not specified in the decorator parameters, the autotuner will use the - default value of (2, 1). Otherwise, the autotuner will use the value specified in the decorator parameters. + :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 instructions. If not specified in the decorator parameters, the autotuner - will use the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters. + :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. If not specified in the decorator parameters, the autotuner will use - the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters. + :param use_tma_store: Whether to use TMA store, defaults to True. :type use_tma_store: bool, optional - :param tolerance: Tolerance value for reference validation comparison, defaults to 1e-01 + :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 + :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 + :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 + :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 + :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 + :param use_tvm_ffi: Whether to use TVM FFI for the kernel, defaults to False. + :type use_tvm_ffi: 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 """ print("Running Blackwell Persistent Dense GEMM test with:") @@ -1781,9 +1874,24 @@ 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'}") + print(f"Use TVM FFI: {'True' if use_tvm_ffi else 'False'}") - # Unpack parameters - m, n, k, l = mnkl + import torch + from cutlass.torch import dtype as torch_dtype + + # Build GEMM object + gemm = PersistentDenseGemmKernel( + acc_dtype, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, use_tma_store + ) + can_implement = gemm.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!") @@ -1793,59 +1901,75 @@ def run( # Get the raw stream pointer as a CUstream current_stream = cuda.CUstream(torch_stream.cuda_stream) - a_tensor, b_tensor, c_tensor, a_torch_cpu, b_torch_cpu, c_torch_cpu, c_torch_gpu = ( - create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype) - ) - - # Build GEMM object - gemm = PersistentDenseGemmKernel( - acc_dtype, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, use_tma_store - ) - # Check if configuration can be implemented - can_implement = gemm.can_implement(a_tensor, b_tensor, c_tensor) - if not can_implement: - raise ValueError( - f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, " - f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, " - f"use_tma_store = {use_tma_store}" - ) max_active_clusters = utils.HardwareInfo().get_max_active_clusters( cluster_shape_mn[0] * cluster_shape_mn[1] ) - compiled_gemm = cute.compile( - gemm, a_tensor, b_tensor, c_tensor, max_active_clusters, current_stream + + options = [] + if use_tvm_ffi: + options.append("--enable-tvm-ffi") + + compiled_fn = compile_bmm( + gemm, + ab_dtype, + ab_dtype, + c_dtype, + a_major, + b_major, + c_major, + max_active_clusters, + current_stream, + options=",".join(options), ) + # Run and verify BMM with torch + a, b, c = prepare_tensors(mnkl, ab_dtype, c_dtype, a_major, b_major, c_major) + if not skip_ref_check: - compiled_gemm(a_tensor, b_tensor, c_tensor, current_stream) - compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance) + # Use small random number for deterministic result for reference check + compiled_fn(a, b, c, torch_stream) + + # Manually quantize to be comparable + ref = ( + torch.bmm(a.to(dtype=torch.float32), b.to(dtype=torch.float32)) + .to(dtype=torch_dtype(c_dtype)) + .to(dtype=torch.float32) + ) + torch.testing.assert_close( + c.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03 + ) + + if not benchmark: + return 0 def generate_tensors(): - a_tensor, _ = cutlass_torch.cute_tensor_like( - a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + init_normal = ab_dtype not in [cutlass.Int8, cutlass.Uint8] + a, b, c = prepare_tensors( + mnkl, + ab_dtype, + c_dtype, + a_major, + b_major, + c_major, + init_random=not init_normal, ) - b_tensor, _ = cutlass_torch.cute_tensor_like( - b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - c_tensor, _ = cutlass_torch.cute_tensor_like( - c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16 - ) - return testing.JitArguments(a_tensor, b_tensor, c_tensor, current_stream) + return testing.JitArguments(a, b, c, torch_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_cpu.numel() * c_torch_cpu.element_size() + a.numel() * a.element_size() + + b.numel() * b.element_size() + + c.numel() * c.element_size() ) workspace_count = testing.get_workspace_count( one_workspace_bytes, warmup_iterations, iterations ) - exec_time = testing.benchmark( - compiled_gemm, + # Return execution time in microseconds + return testing.benchmark( + compiled_fn, workspace_generator=generate_tensors, workspace_count=workspace_count, stream=current_stream, @@ -1853,18 +1977,17 @@ def run( iterations=iterations, ) - return exec_time # Return execution time in microseconds + +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." + ) -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 prepare_parser(): parser = argparse.ArgumentParser( description="Example of Dense Persistent GEMM on Blackwell." @@ -1872,19 +1995,13 @@ if __name__ == "__main__": parser.add_argument( "--mnkl", - type=parse_comma_separated_ints, + 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, + type=_parse_comma_separated_ints, default=(1, 1), help="Cluster shape (comma-separated)", ) @@ -1905,6 +2022,9 @@ if __name__ == "__main__": parser.add_argument( "--tolerance", type=float, default=1e-01, help="Tolerance for validation" ) + parser.add_argument( + "--benchmark", action="store_true", help="Only benchmark the kernel" + ) parser.add_argument( "--warmup_iterations", type=int, default=0, help="Warmup iterations" ) @@ -1923,6 +2043,24 @@ if __name__ == "__main__": default=False, help="Use circular buffer tensor sets to ensure L2 cold cache", ) + parser.add_argument( + "--use_tvm_ffi", + action="store_true", + default=False, + help="Enable TVM FFI for the kernel, defaults to False using CuTe DSL's native runtime", + ) + + 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() @@ -1952,5 +2090,7 @@ if __name__ == "__main__": args.iterations, args.skip_ref_check, args.use_cold_l2, + args.use_tvm_ffi, + args.benchmark, ) print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py b/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py index 3876c288..70e8b1e2 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py @@ -36,6 +36,7 @@ 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 @@ -535,6 +536,7 @@ class DenseGemmKernel: 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 @@ -549,6 +551,7 @@ class DenseGemmKernel: producer_group=acc_pipeline_producer_group, consumer_group=acc_pipeline_consumer_group, cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, ) acc_producer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Producer, self.num_acc_stage @@ -569,8 +572,7 @@ class DenseGemmKernel: ) # Cluster arrive after barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # # Setup smem tensor A/B/C @@ -686,8 +688,7 @@ class DenseGemmKernel: # # Cluster wait before tensor memory alloc # - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_wait() + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) # Alloc tensor memory buffer tmem.allocate(self.num_tmem_alloc_cols) diff --git a/examples/python/CuTeDSL/blackwell/fmha.py b/examples/python/CuTeDSL/blackwell/fmha.py index 58905eef..d1e9c927 100644 --- a/examples/python/CuTeDSL/blackwell/fmha.py +++ b/examples/python/CuTeDSL/blackwell/fmha.py @@ -48,9 +48,11 @@ import cutlass.cute.testing as testing from cutlass.cute.runtime import from_dlpack from cutlass.cute.typing import Int32, Int64, Float32 -current_dir = os.path.dirname(os.path.abspath(__file__)) -sys.path.insert(0, os.path.join(current_dir, "..")) -from utils import fmha_helpers as fmha_utils +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 """ A fused multi-head attention (FMHA) example for the NVIDIA Blackwell SM100 architecture using CUTE DSL diff --git a/examples/python/CuTeDSL/blackwell/fmha_bwd.py b/examples/python/CuTeDSL/blackwell/fmha_bwd.py index 88820c1b..2cfe72a9 100644 --- a/examples/python/CuTeDSL/blackwell/fmha_bwd.py +++ b/examples/python/CuTeDSL/blackwell/fmha_bwd.py @@ -49,9 +49,11 @@ import cutlass.utils.blackwell_helpers as sm100_utils from cutlass.cute.runtime import from_dlpack from cutlass.cute.typing import Int32, Float32, Float8E4M3FN, Float16, BFloat16, Boolean -current_dir = os.path.dirname(os.path.abspath(__file__)) -sys.path.insert(0, os.path.join(current_dir, "..")) -from utils import fmha_helpers as fmha_utils +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 """ A fused multi-head attention (FMHA) backward pass example for the NVIDIA Blackwell SM100 architecture using CUTE DSL diff --git a/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py b/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py index e8001903..53cec9d9 100644 --- a/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py +++ b/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py @@ -40,6 +40,7 @@ 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 import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.utils.blockscaled_layout as blockscaled_utils from cutlass.cute.runtime import from_dlpack @@ -177,22 +178,18 @@ class Sm100GroupedBlockScaledGemmKernel: self.threads_per_cta = 32 * len( (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id) ) - # Set barrier for cta sync, epilogue sync and tmem ptr sync - self.cta_sync_barrier = pipeline.NamedBarrier( - barrier_id=1, - num_threads=self.threads_per_cta, - ) + # Set barrier for epilogue sync and tmem ptr sync self.epilog_sync_barrier = pipeline.NamedBarrier( - barrier_id=2, + barrier_id=1, num_threads=32 * len(self.epilog_warp_id), ) self.tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=3, + barrier_id=2, num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)), ) # Barrier used by MMA/TMA warps to signal A/B tensormap initialization completion self.tensormap_ab_init_barrier = pipeline.NamedBarrier( - barrier_id=4, + barrier_id=3, num_threads=64, ) self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") @@ -646,6 +643,7 @@ class Sm100GroupedBlockScaledGemmKernel: cluster=(*self.cluster_shape_mn, 1), smem=self.shared_storage.size_in_bytes(), stream=stream, + min_blocks_per_mp=1, ) return @@ -781,11 +779,9 @@ class Sm100GroupedBlockScaledGemmKernel: cute.arch.mbarrier_init( tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads ) - cute.arch.mbarrier_init_fence() # Cluster arrive after barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # # Setup smem tensor A/B/SFA/SFB/C @@ -944,10 +940,7 @@ class Sm100GroupedBlockScaledGemmKernel: # # Cluster wait before tensor memory alloc # - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_wait() - else: - self.cta_sync_barrier.arrive_and_wait() + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) # # Get tensormap buffer address @@ -2894,6 +2887,7 @@ def run( tensor_of_tensormap, max_active_clusters, current_stream, + options=f"--opt-level 2", ) # reference check diff --git a/examples/python/CuTeDSL/blackwell/grouped_gemm.py b/examples/python/CuTeDSL/blackwell/grouped_gemm.py index 50d455da..becac750 100644 --- a/examples/python/CuTeDSL/blackwell/grouped_gemm.py +++ b/examples/python/CuTeDSL/blackwell/grouped_gemm.py @@ -39,6 +39,7 @@ 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 from cutlass.cute.nvgpu import cpasync, tcgen05 import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.torch as cutlass_torch @@ -153,22 +154,18 @@ class GroupedGemmKernel: self.threads_per_cta = 32 * len( (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id) ) - # Set barrier for cta sync, epilog sync, tmem ptr sync and tensormap update sync - self.cta_sync_barrier = pipeline.NamedBarrier( - barrier_id=1, - num_threads=self.threads_per_cta, - ) + # Set barrier for epilog sync, tmem ptr sync and tensormap update sync self.epilog_sync_barrier = pipeline.NamedBarrier( - barrier_id=2, + barrier_id=1, num_threads=32 * len(self.epilog_warp_id), ) self.tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=3, + barrier_id=2, num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)), ) # Barrier used by MMA/TMA warps to signal A/B tensormap initialization completion self.tensormap_ab_init_barrier = pipeline.NamedBarrier( - barrier_id=4, + barrier_id=3, num_threads=32 * (len(self.epilog_warp_id) + 1), ) self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") @@ -586,11 +583,9 @@ class GroupedGemmKernel: is_two_cta=use_2cta_instrs, two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, ) - cute.arch.mbarrier_init_fence() # Cluster arrive after barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # # Setup smem tensor A/B/C @@ -718,10 +713,7 @@ class GroupedGemmKernel: # # Cluster wait before tensor memory alloc # - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_wait() - else: - self.cta_sync_barrier.arrive_and_wait() + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) # # Get tensormap buffer address diff --git a/examples/python/CuTeDSL/blackwell/grouped_mixed_input_gemm.py b/examples/python/CuTeDSL/blackwell/grouped_mixed_input_gemm.py new file mode 100644 index 00000000..087c8b47 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/grouped_mixed_input_gemm.py @@ -0,0 +1,3202 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import argparse +from math import log2, ceil +from typing import Optional, Union + +import torch +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import ( + extract_mlir_values, + new_from_mlir_values, +) +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.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 +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 could 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 explititly 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 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/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/grouped_grouped_mixed_input \ + --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 ContiguousGGSearchState: + """ + The state of group search for grouped gemm with contiguous offsets. + + The state records the progress of group seach algorithm on 1 mode. It will be + initialized once and updated in every round of group index search. + + :param last_tile_count: Number of cluster tiles before the current group + :type last_tile_count: cutlass.Int32 + :param cur_boundary: The boundary of the current group, which is the size along the seach + mode before the next group + :type cur_boundary: cutlass.Int32 + :param cur_tile_count: Number of cluster tiles searched so far + :type cur_tile_count: cutlass.Int32 + :param cur_group_idx: The index of the current group + :type cur_group_idx: cutlass.Int32 + :param cur_offset: The starting offset of the current group along the search mode + :type cur_offset: cutlass.Int32 + :param cur_start: The starting offset of the current cluster tile size along the search mode + when group search is done + :type cur_start: cutlass.Int32 + """ + + def __init__( + self, + last_tile_count: cutlass.Int32, + cur_boundary: cutlass.Int32, + cur_tile_count: cutlass.Int32, + cur_group_idx: cutlass.Int32, + cur_offset: cutlass.Int32, + cur_start: cutlass.Int32, + ): + self.last_tile_count = last_tile_count + self.cur_boundary = cur_boundary + self.cur_tile_count = cur_tile_count + self.cur_group_idx = cur_group_idx + self.cur_offset = cur_offset + self.cur_start = cur_start + + def __extract_mlir_values__(self): + values = extract_mlir_values(self.last_tile_count) + values.extend(extract_mlir_values(self.cur_boundary)) + values.extend(extract_mlir_values(self.cur_tile_count)) + values.extend(extract_mlir_values(self.cur_group_idx)) + values.extend(extract_mlir_values(self.cur_offset)) + values.extend(extract_mlir_values(self.cur_start)) + return values + + def __new_from_mlir_values__(self, values) -> "ContiguousGGSearchState": + last_tile_count = new_from_mlir_values(self.last_tile_count, [values[0]]) + cur_boundary = new_from_mlir_values(self.cur_boundary, [values[1]]) + cur_tile_count = new_from_mlir_values(self.cur_tile_count, [values[2]]) + cur_group_idx = new_from_mlir_values(self.cur_group_idx, [values[3]]) + cur_offset = new_from_mlir_values(self.cur_offset, [values[4]]) + cur_start = new_from_mlir_values(self.cur_start, [values[5]]) + return ContiguousGGSearchState( + last_tile_count, + cur_boundary, + cur_tile_count, + cur_group_idx, + cur_offset, + cur_start, + ) + + +def create_initial_search_state() -> ContiguousGGSearchState: + """ + Create an initial search state for grouped gemm with contiguous offsets. + """ + return ContiguousGGSearchState( + last_tile_count=cutlass.Int32(0), + cur_boundary=cutlass.Int32(0), + cur_tile_count=cutlass.Int32(0), + cur_group_idx=cutlass.Int32(0), + cur_offset=cutlass.Int32(0), + cur_start=cutlass.Int32(0), + ) + + +class GroupedWorkTileInfo: + """ + Tile info for grouped gemm with contiguous offsets. + It's consutrcted from the search state and contains informtion needed for different warps. + + :param group_count: The total number of groups + :type group_count: int + :param cta_coord_m: The coordinate of the current CTA tile along the M mode + :type cta_coord_m: cutlass.Int32 + :param coord_n: The starting offset on N mode for the current CTA tile + :type coord_n: cutlass.Int32 + :param group_idx: The index of the current group + :type group_idx: cutlass.Int32 + :param distance_to_boundary: The distance to the boundary of the current group + :type distance_to_boundary: cutlass.Int32 + """ + + def __init__( + self, + group_count: int, + cta_coord_m: cutlass.Int32, + coord_n: cutlass.Int32, + group_idx: cutlass.Int32, + distance_to_boundary: cutlass.Int32, + ): + self.cta_coord_m = cta_coord_m + self.coord_n = coord_n + self.group_idx = group_idx + self.distance_to_boundary = distance_to_boundary + self.group_count = group_count + + def __extract_mlir_values__(self): + values = extract_mlir_values(self.cta_coord_m) + values.extend(extract_mlir_values(self.coord_n)) + values.extend(extract_mlir_values(self.group_idx)) + values.extend(extract_mlir_values(self.distance_to_boundary)) + return values + + def __new_from_mlir_values__(self, values): + assert len(values) == 4 + new_cta_coord_m = new_from_mlir_values(self.cta_coord_m, [values[0]]) + new_coord_n = new_from_mlir_values(self.coord_n, [values[1]]) + new_group_idx = new_from_mlir_values(self.group_idx, [values[2]]) + new_distance_to_boundary = new_from_mlir_values( + self.distance_to_boundary, [values[3]] + ) + return GroupedWorkTileInfo( + self.group_count, + new_cta_coord_m, + new_coord_n, + new_group_idx, + new_distance_to_boundary, + ) + + @property + def is_valid_tile(self): + return self.group_idx < self.group_count + + +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 group seach algorithm + is applied on the N mode to find the group index for each CTA tile. A cumsum input tensor provides the offset 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 + """ + + 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, + ): + """ + 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.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 subtile + - 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 = GroupedMixedInputGemmKernel.align_up( + self.num_acc_tmem_cols + self.num_a_tmem_cols, + cute.arch.SM100_TMEM_MIN_ALLOC_COLUMNS, + ) + 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 should 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, + ) + 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: + # 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 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) + ) + 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() + # 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.warpgroup_reg_dealloc(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 = create_initial_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 = self.group_search( + 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( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_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.warpgroup_reg_dealloc(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 = self.make_work_tile_info( + sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_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 = self.make_work_tile_info( + sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_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.warpgroup_reg_dealloc(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 = self.make_work_tile_info( + sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_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 = self.make_work_tile_info( + sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_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.warpgroup_reg_alloc(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 = 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 subtile 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 = self.make_work_tile_info( + sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_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 = ( + tArA_load[(None, idx)].load().to(self.mma_dtype) + ) + 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 + if cutlass.const_expr(dst_copy_a is not None): + cute.copy( + dst_copy_a, + tArA_transform, + tAsA_transform[ + (None, None, None, None, trans2mma_producer_state.index) + ], + ) + else: + cute.autovec_copy( + tArA_transform, + tAsA_transform[ + (None, None, None, None, trans2mma_producer_state.index) + ], + ) + # 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( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_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 = self.make_work_tile_info( + sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_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.warpgroup_reg_dealloc(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 = self.make_work_tile_info( + sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_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 = self.make_work_tile_info( + sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_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.warpgroup_reg_alloc(self.num_regs_epilogue_warps) + epi_tidx = tidx + tCtAcc_base = accumulators + # Partition for epilogue + 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, simt_atom, tTR_gC_partitioned) = ( + self.epilog_gmem_copy_and_partition( + 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 = self.make_work_tile_info( + sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_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( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_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 = self.make_work_tile_info( + sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_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() + + @cute.jit + def group_search( + self, + group_count: cutlass.Int32, + linear_idx: cutlass.Int32, + search_state: ContiguousGGSearchState, + cumsum: cute.Tensor, + search_mode: int, + ) -> ContiguousGGSearchState: + """ + Group search for contiguously grouped gemm. + """ + not_found = linear_idx >= search_state.cur_tile_count + next_boundary = cutlass.Int32(0) + cur_group_idx = search_state.cur_group_idx + cur_offset = search_state.cur_offset + last_tile_count = search_state.last_tile_count + cur_boundary = search_state.cur_boundary + cur_tile_count = search_state.cur_tile_count + if not_found: + cur_group_idx = cur_group_idx + 1 + while not_found and cur_group_idx <= group_count: + next_boundary = cumsum[cur_group_idx] + num_m_blocks = cute.ceil_div( + (next_boundary - cur_boundary), + self.cluster_tile_shape_mnk[search_mode], + ) + next_tile_count = num_m_blocks + cur_tile_count + not_found = linear_idx >= next_tile_count + + last_tile_count = cur_tile_count + cur_offset = cur_boundary + cur_boundary = next_boundary + cur_tile_count = next_tile_count + if not_found: + cur_group_idx = cur_group_idx + 1 + cur_start = cur_offset + self.cluster_tile_shape_mnk[search_mode] * ( + linear_idx - last_tile_count + ) + return ContiguousGGSearchState( + last_tile_count, + cur_boundary, + cur_tile_count, + cur_group_idx, + cur_offset, + cur_start, + ) + + def make_work_tile_info(self, sTile_info: cute.Tensor): + tile_info = cute.make_rmem_tensor(sTile_info.shape, sTile_info.element_type) + cute.autovec_copy(sTile_info, tile_info) + return GroupedWorkTileInfo( + self.group_count, tile_info[0], tile_info[1], tile_info[2], tile_info[3] + ) + + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tma_atom_c: cute.CopyAtom, + tiled_copy_t2r: cute.TiledCopy, + gC_mnl_tma: cute.Tensor, + gC_mnl_simt: cute.Tensor, + epi_tile: cute.Tile, + sC: cute.Tensor, + ) -> tuple[cute.CopyAtom, cute.Tensor, cute.Tensor, cute.CopyAtom, cute.Tensor]: + """ + Partitions source and destination tensors for a global memory store. + + This method generates a tiled copy for storing results to global memory + and partitions the source (register or shared memory) and destination + (global memory) tensors accordingly. The behavior varies based on whether + TMA store is enabled. + + :param tidx: The thread index in epilogue warp groups. + :type tidx: cutlass.Int32 + :param tma_atom_c: The TMA copy atom. + :type tma_atom_c: cute.CopyAtom + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy. + :type tiled_copy_t2r: cute.TiledCopy + :param gC_mnl_tma: The global tensor C for TMA. + :type gC_mnl_tma: cute.Tensor + :param gC_mnl_simt: The global tensor C for SIMT Copy. + :type gC_mnl_simt: cute.Tensor + :param epi_tile: The epilogue tiler. + :type epi_tile: cute.Tile + :param sC: The shared memory tensor C. + :return: A tuple containing the appropriate copy atom and partitioned + source and destination tensors for the store operation. + :rtype: tuple[cute.CopyAtom, cute.Tensor, cute.Tensor, cute.CopyAtom, cute.Tensor] + """ + gC_epi_tma = cute.flat_divide( + gC_mnl_tma[((None, None), 0, 0, None, None, None)], epi_tile + ) + gC_epi_simt = cute.flat_divide( + gC_mnl_simt[((None, None), 0, 0, None, None, None)], epi_tile + ) + # TMA store + sC_for_tma_partition = cute.group_modes(sC, 0, 2) + gC_for_tma_partition = cute.group_modes(gC_epi_tma, 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, + ) + # SIMT Store + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_gC = thr_copy_t2r.partition_D(gC_epi_simt) + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + return tma_atom_c, bSG_sC, bSG_gC, simt_atom, tTR_gC + + 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]: + """ + Partitions source and destination tensors for a shared memory store. + + This method generates a tiled copy for storing results to shared memory + and partitions the source (register) and destination (shared memory) + tensors accordingly. + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy. + :param tTR_rC: The partitioned accumulator tensor. + :param tidx: The thread index in epilogue warp groups. + :param sC: The shared memory tensor to be copied and partitioned. + :return: A tuple containing the tiled copy for the store operation and + the partitioned source and destination tensors. + """ + 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_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]: + """ + Partitions source and destination tensors for a tensor memory load. + + This method generates a tiled copy for loading accumulators from tensor + memory and partitions the source (tensor memory) and destination + (register) tensors accordingly. + + :param tidx: The thread index in epilogue warp groups. + :param tAcc: The accumulator tensor to be copied and partitioned. + :param gC_mnl: The global tensor C. + :param epi_tile: The epilogue tiler. + :param use_2cta_instrs: Whether use_2cta_instrs is enabled. + :return: A tuple containing the tiled copy for the load operation and + the partitioned source and destination tensors. + """ + # 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 + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL) + 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 + + @staticmethod + def align_up(x: int, align: int) -> int: + """Align x up to the nearest multiple of align.""" + return (x + align - 1) // align * align + + @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.SM100_TMEM_CAPACITY_COLUMNS + 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 = GroupedMixedInputGemmKernel.align_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 = GroupedMixedInputGemmKernel.align_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 + caveout_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 = GroupedMixedInputGemmKernel.align_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 = GroupedMixedInputGemmKernel.align_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 = GroupedMixedInputGemmKernel.align_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 = ( + GroupedMixedInputGemmKernel.align_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 - caveout_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 + - caveout_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 is_valid_tensor_alignment( + m: int, + n: int, + k: int, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + scale_dtype: type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mnk: tuple[int, int, int], + use_2cta_instrs: bool, + cluster_shape_mn: tuple[int, int], + scale_granularity_m: int, + scale_granularity_k: int, + ) -> bool: + """ + Check if the tensor alignments are valid for the given problem size and data types. + """ + + 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)) + and check_contiguous_16B_alignment(b_dtype, b_major == "n", (n, k)) + and check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n)) + and ( + scale_granularity_k == 0 + or check_contiguous_16B_alignment( + b_dtype, True, (m, k // scale_granularity_k) + ) + ) + ): + return False + # Check if scale tensor matches the TMA load 128B alignment requirement + cta_tile_shape_mnk = ( + mma_tiler_mnk[0] // (2 if use_2cta_instrs else 1), + mma_tiler_mnk[1], + mma_tiler_mnk[2], + ) + if ( + scale_granularity_m > 0 + and (cta_tile_shape_mnk[0] // cluster_shape_mn[1] // scale_granularity_m) + * (scale_dtype.width // 8) + < 128 + ): + return False + + return True + + def is_valid_mma_tiler_and_cluster_shape( + mma_tiler: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + ) -> bool: + """ + Check if the MMA tiler and cluster shape are valid for the given problem size. + """ + if cluster_shape_mn[0] % (2 if use_2cta_instrs else 1) != 0: + return False + + if (mma_tiler[0] // (2 if use_2cta_instrs else 1)) not in [64, 128]: + 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, + ) -> bool: + """ + Check if the kernel can be implemented for the given tensor shapes and data types. + """ + m, n, k, l = mnkl + + if not GroupedMixedInputGemmKernel.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 GroupedMixedInputGemmKernel.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 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], + 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 torch.maximum(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") + # 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], + 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, + 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( + 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], + scale_granularity_m: int = 0, + scale_granularity_k: int = 0, + uniform_group_sizes: bool = False, +) -> tuple: + """ + Create all input and output tensors for the 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, + 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 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( + 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 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 + mixed_input_gemm = GroupedMixedInputGemmKernel( + scale_granularity_m, + scale_granularity_k, + acc_dtype, + use_2cta_instrs, + mma_tiler_mnk, + cluster_shape_mn, + group_count, + ) + ( + 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, + 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], + ) + compiled_kernel = cute.compile( + mixed_input_gemm, + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + max_active_clusters, + current_stream, + ) + + if not skip_ref_check: + compiled_kernel( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + current_stream, + ) + 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, a_torch_cpu, a_scale_torch_cpu = create_tensor_a( + l, + m, + k, + a_major, + a_dtype, + scale_granularity_m, + scale_granularity_k, + b_dtype, + ) + num_groups = l + fused_n = n * num_groups + b_torch_cpu = cutlass_torch.matrix( + 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, 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), + ) + alignment_n = 16 * 8 // b_dtype.width + cumsum_tensor, cumsum_torch_cpu = create_cumsum_tensor( + num_groups, fused_n, alignment_n, uniform_distribution=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/mamba2_ssd/mamba2_ssd.py b/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py index afa8830b..b82b08bf 100644 --- a/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py +++ b/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py @@ -27,6 +27,8 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +import os +import sys import argparse from typing import List, Type, Tuple, Optional import cuda.bindings.driver as cuda @@ -39,21 +41,22 @@ 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 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 -import sys -from pathlib import Path +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) -sys.path.append(str(Path(__file__).resolve().parent)) -from mamba2_ssd_reference import ( +from blackwell.mamba2_ssd.mamba2_ssd_reference import ( ssd_reference_fp32_all, ssd_reference_lowprecision_intermediates, analyze_relative_diffs, ) -from mamba2_ssd_tile_scheduler import ( +from blackwell.mamba2_ssd.mamba2_ssd_tile_scheduler import ( Mamba2SSDTileSchedulerParams, Mamba2SSDTileScheduler, ) @@ -811,12 +814,10 @@ class SSDKernel: ) # Cluster arrive after barrier init - if cute.size(self.cluster_shape_mnk) > 1: - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True) # Cluster wait before tmem alloc - if cute.size(self.cluster_shape_mnk) > 1: - cute.arch.cluster_wait() + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk) tmem_alloc_barrier = pipeline.NamedBarrier( barrier_id=0, @@ -2574,6 +2575,7 @@ class SSDKernel: consumer_group=x_consumer_group, tx_count=self.num_x_load_bytes, barrier_storage=x_full_mbar_ptr, + defer_sync=True, ) else: x_consumer_group_umma = pipeline.CooperativeGroup( @@ -2590,6 +2592,7 @@ class SSDKernel: consumer_group_async=x_consumer_group_async, tx_count=self.num_x_load_bytes, barrier_storage=x_full_mbar_ptr, + defer_sync=True, ) def make_and_init_b_pipeline(self, b_full_mbar_ptr): @@ -2609,6 +2612,7 @@ class SSDKernel: consumer_group_async=b_consumer_group_async, tx_count=self.num_b_load_bytes, barrier_storage=b_full_mbar_ptr, + defer_sync=True, ) def make_and_init_c_pipeline(self, c_full_mbar_ptr): @@ -2624,6 +2628,7 @@ class SSDKernel: consumer_group=c_consumer_group, tx_count=self.num_c_load_bytes, barrier_storage=c_full_mbar_ptr, + defer_sync=True, ) def make_and_init_deltas_pipeline(self, deltas_full_mbar_ptr): @@ -2643,6 +2648,7 @@ class SSDKernel: consumer_group=deltas_consumer_group, tx_count=self.num_delta_load_bytes + self.num_cumsum_delta_load_bytes, barrier_storage=deltas_full_mbar_ptr, + defer_sync=True, ) def make_and_init_d_pipeline(self, d_full_mbar_ptr): @@ -2662,6 +2668,7 @@ class SSDKernel: consumer_group=d_consumer_group, tx_count=self.num_d_load_bytes, barrier_storage=d_full_mbar_ptr, + defer_sync=True, ) def make_and_init_intra1_acc_pipeline(self, intra1_acc_full_mbar_ptr): @@ -2676,6 +2683,7 @@ class SSDKernel: producer_group=intra1_acc_producer_group, consumer_group=intra1_acc_consumer_group, barrier_storage=intra1_acc_full_mbar_ptr, + defer_sync=True, ) def make_and_init_intra2_q_pipeline(self, intra2_q_full_mbar_ptr): @@ -2690,6 +2698,7 @@ class SSDKernel: producer_group=intra2_q_producer_group, consumer_group=intra2_q_consumer_group, barrier_storage=intra2_q_full_mbar_ptr, + defer_sync=True, ) def make_and_init_intra2_acc_pipeline(self, intra2_acc_full_mbar_ptr): @@ -2704,6 +2713,7 @@ class SSDKernel: producer_group=intra2_acc_producer_group, consumer_group=intra2_acc_consumer_group, barrier_storage=intra2_acc_full_mbar_ptr, + defer_sync=True, ) def make_and_init_inter1_b_pipeline(self, inter1_b_full_mbar_ptr): @@ -2718,6 +2728,7 @@ class SSDKernel: producer_group=inter1_b_producer_group, consumer_group=inter1_b_consumer_group, barrier_storage=inter1_b_full_mbar_ptr, + defer_sync=True, ) def make_and_init_inter1_acc_pipeline(self, inter1_acc_full_mbar_ptr): @@ -2732,6 +2743,7 @@ class SSDKernel: producer_group=inter1_acc_producer_group, consumer_group=inter1_acc_consumer_group, barrier_storage=inter1_acc_full_mbar_ptr, + defer_sync=True, ) def make_and_init_inter2_p_pipeline(self, inter2_p_full_mbar_ptr): @@ -2746,6 +2758,7 @@ class SSDKernel: producer_group=inter2_p_producer_group, consumer_group=inter2_p_consumer_group, barrier_storage=inter2_p_full_mbar_ptr, + defer_sync=True, ) def make_and_init_inter2_acc_pipeline(self, inter2_acc_full_mbar_ptr): @@ -2760,6 +2773,7 @@ class SSDKernel: producer_group=inter2_acc_producer_group, consumer_group=inter2_acc_consumer_group, barrier_storage=inter2_acc_full_mbar_ptr, + defer_sync=True, ) def tma_partition_for_mma_b_operand( diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill.py b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill.py new file mode 100644 index 00000000..8483e78d --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill.py @@ -0,0 +1,2353 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import argparse +import math +import os +import sys +from dataclasses import dataclass +from typing import Type, Tuple, Optional + +import torch +import cuda.bindings.driver as cuda + +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 ( + CooperativeGroup, + PipelineOp, + PipelineState, + pipeline_init_wait, + PipelineAsync, +) +import cutlass.torch as cutlass_torch +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.typing import Int32, Int64, Float32, Boolean +from cutlass.cutlass_dsl import if_generate + +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 + + +def make_thread_cooperative_group(size: int): + return pipeline.CooperativeGroup(pipeline.Agent.Thread, size) + + +@dataclass(frozen=True) +class PipelineTmaTransform(PipelineAsync): + @staticmethod + def create( + *, + num_stages: int, + producer_group: CooperativeGroup, + consumer_group: CooperativeGroup, + tx_count: int, + barrier_storage: cute.Pointer = None, + cta_layout_vmnk: Optional[cute.Layout] = None, + ): + if not isinstance(barrier_storage, cute.Pointer): + raise ValueError( + f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" + ) + + producer_type = PipelineOp.TmaLoad + consumer_type = PipelineOp.AsyncThread + + producer = (producer_type, producer_group) + consumer = (consumer_type, consumer_group) + + sync_object_full = PipelineAsync._make_sync_object( + barrier_storage.align(min_align=8), num_stages, producer, tx_count + ) + sync_object_empty = PipelineAsync._make_sync_object( + barrier_storage.align(min_align=8) + num_stages, num_stages, consumer + ) + + pipeline_init_wait() + + return PipelineTmaTransform( + sync_object_full, + sync_object_empty, + num_stages, + producer_mask=None, + consumer_mask=None, + ) + + def producer_acquire( + self, state: PipelineState, try_acquire_token: Optional[Boolean] = None + ): + """ + TMA producer commit conditionally waits on buffer empty and sets the transaction barrier. + """ + if_generate( + try_acquire_token is None or try_acquire_token == 0, + lambda: self.sync_object_empty.wait(state.index, state.phase), + ) + self.sync_object_full.arrive(state.index, self.producer_mask) + + def producer_commit(self, state: PipelineState): + """ + TMA producer commit is a noop since TMA instruction itself updates the transaction count. + """ + pass + + +class MixedInputFusedMultiHeadAttentionPrefill: + def __init__( + self, + scale_granularity: int, + qk_acc_dtype: Type[cutlass.Numeric], + pv_acc_dtype: Type[cutlass.Numeric], + cta_tiler: Tuple[int, int, int], # seq_q, seq_k, d + is_persistent: bool, + mask_type: fmha_utils.MaskEnum, + ): + self.qk_acc_dtype = qk_acc_dtype + self.pv_acc_dtype = pv_acc_dtype + + self.qk_mma_tiler = ( + cta_tiler[0] * 2, # default 2cta + cta_tiler[1], # GemmN at most 256 + min(cta_tiler[2], 64), # 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.cta_tiler = cta_tiler + self.scale_granularity = scale_granularity + + self.iterations_qk = cta_tiler[2] // self.qk_mma_tiler[2] + self.iterations_pv_k = cta_tiler[1] // self.pv_mma_tiler[2] + self.iterations_pv_n = cta_tiler[2] // self.pv_mma_tiler[1] + self.iterations_pv = self.iterations_pv_k * self.iterations_pv_n + self.cluster_shape_mn = (2, 1) # use 2x1 cluster by default + self.tmem_warp_shape_mn = (2, 2) + 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 + + SM100_TMEM_CAPACITY_COLUMNS = 512 + self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + + self.tmem_alloc_sync_bar_id = 1 + self.tmem_s_offset = 0 + self.tmem_p_offset = self.tmem_s_offset + self.tmem_o_offset = 256 + self.num_regs_softmax = 192 + self.num_regs_other = 96 + self.num_regs_transform = 112 + self.buffer_align_bytes = 1024 + self.threads_per_warp = 32 + self.smem_exchange_sync_bar = pipeline.NamedBarrier( + barrier_id=2, + num_threads=(self.threads_per_warp * len(self.softmax_warp_ids)), + ) + 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 = 5 + 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 = 3 + + @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_m = cute.make_layout(self.pv_block_tiler[0]) + self.epi_tile_n = cute.make_layout( + (32, 2), stride=(1, self.pv_block_tiler[1] // self.tmem_warp_shape_mn[1]) + ) + self.epi_tile = (self.epi_tile_m, self.epi_tile_n) + + 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.k_dtype, + self.kv_stage, + ) + 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.iterations_pv_k * self.qk_acc_stage), + ) + p_smem_layout_staged = cute.logical_divide( + p_smem_layout_staged, (None, None, None, self.iterations_pv_k) + ) + 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 = ( + self.get_scale_smem_layout( + 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 = ( + self.get_scale_smem_layout( + self.pv_mma_tiler, + self.v_major_mode, + ) + ) + scale_v_smem_layout_staged = cute.append( + scale_v_smem_layout, + cute.make_layout( + (self.iterations_pv_k, self.scale_v_stage), + stride=( + cute.cosize(scale_v_smem_layout.outer), + cute.cosize(scale_v_smem_layout.outer) * self.iterations_pv_k, + ), + ), + ) + scale_v_s2r_view_layout_staged = cute.append( + scale_v_s2r_view_layout, + cute.make_layout( + (self.iterations_pv_k, self.scale_v_stage), + stride=( + cute.cosize(scale_v_s2r_view_layout), + cute.cosize(scale_v_s2r_view_layout) * self.iterations_pv_k, + ), + ), + ) + + 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) + * self.iterations_pv_k + ) + + @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] + tmem_dealloc_mbar_ptr: Int64 + tmem_holding_buf: Int32 + + self.shared_storage = SharedStorage + + # 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=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(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, + ).make_participants() + load_kv_producer, load_kv_consumer = PipelineTmaTransform.create( + num_stages=self.kv_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group( + 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(), + ).make_participants() + load_scale_k_producer, load_scale_k_consumer = PipelineTmaTransform.create( + num_stages=self.scale_k_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group( + 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(), + ).make_participants() + load_scale_v_producer, load_scale_v_consumer = PipelineTmaTransform.create( + num_stages=self.scale_v_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group( + 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(), + ).make_participants() + dequant_kv_producer, dequant_kv_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.kv_trans_stage, + producer_group=make_thread_cooperative_group( + len(self.transform_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0] + ), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + barrier_storage=storage.dequant_kv_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + ).make_participants() + mma_s_producer, mma_s_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.qk_acc_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group( + 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, + ).make_participants() + p_mma_producer, p_mma_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.qk_acc_stage, + producer_group=make_thread_cooperative_group( + len(self.softmax_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0] + ), + consumer_group=make_thread_cooperative_group(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=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group( + 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, + ).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 + cute.arch.cluster_arrive_relaxed() + sQ = smem.allocate_tensor( + element_type=self.q_dtype, + layout=q_smem_layout_staged.outer, + swizzle=q_smem_layout_staged.inner, + byte_alignment=128, + ) + 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) + 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, + ) + smem_exchange = 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, + ) + 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_n, + 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 = self.scale_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 = self.scale_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() + cute.arch.cluster_wait() + + # /////////////////////////////////////////////////////////////////////////////// + # Load + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.load_warp_id: + cute.arch.warpgroup_reg_dealloc(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 = ( + self.load_qk( # Q & K0 & ScaleK0 + 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 = ( + self.load_qk( # Ki & ScaleKi + 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 = ( + self.load_v( # Vi-1 & ScaleVi-1 + 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 = ( + self.load_v( # Vend & ScaleVend + 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.warpgroup_reg_dealloc(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 = ( + self.mma_qk( # QK0 + 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 = self.mma_qk( # QKi + 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, + (tOtO_staged, tOrP, tOrV_trans), + (p_mma_consumer, mma_o_producer, dequant_kv_consumer), + ) + mma_s_producer, _, dequant_kv_consumer = ( + self.mma_qk( # QKend needs to release Q + 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, + (tOtO_staged, tOrP, tOrV_trans), + (p_mma_consumer, mma_o_producer, dequant_kv_consumer), + ) + else: + mma_s_producer, load_q_consumer, dequant_kv_consumer = ( + self.mma_qk( # QK0 + 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, + (tOtO_staged, tOrP, tOrV_trans), + (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.mma_warp_id and warp_idx >= self.softmax_warp_ids[0]: + cute.arch.warpgroup_reg_alloc(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, smem_exchange), + (mma_s_consumer, p_mma_producer), + ) + row_max_prev = row_max + 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 + row_max, row_sum, mma_s_consumer, p_mma_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, sP, smem_exchange), + (mma_s_consumer, p_mma_producer), + ) + ) + # Oi-1 -> Oi + mma_o_consumer = self.correction_rescale( + (row_max, row_max_prev, scale_softmax_log2), + (mma_o_consumer, tOtO_staged, cO_staged), + epi_tile, + ) + row_max_prev = row_max + # O_partial -> O_final + mma_o_consumer = self.correction_epilog( + (row_sum, seqlen_q, scale_output), + (mma_o_consumer, gO_staged, cO_staged, tOtO_staged, smem_exchange), + epi_tile, + ) + 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.warpgroup_reg_dealloc(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 = ( + self.dequant_k( # K0 + (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 = ( + self.dequant_k( # Ki + (sK, sScaleK_, sK_trans), + ( + load_kv_consumer, + load_scale_k_consumer, + dequant_kv_producer, + ), + ) + ) + load_kv_consumer, load_scale_v_consumer, dequant_kv_producer = ( + self.dequant_v( # Vi-1 + (sV, sScaleV_, sV_trans), + ( + load_kv_consumer, + load_scale_v_consumer, + dequant_kv_producer, + ), + mma_tile_coord_v, + ) + ) + load_kv_consumer, load_scale_v_consumer, dequant_kv_producer = ( + self.dequant_v( # Vend + (sV, sScaleV_, sV_trans), + (load_kv_consumer, load_scale_v_consumer, dequant_kv_producer), + mma_tile_coord_v, + ) + ) + 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.warpgroup_reg_dealloc(self.num_regs_other) + + return + + @cute.jit + def load_qk( + self, + 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(self.iterations_qk, unroll_full=True): + 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( + self, + 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() + for iter_k in cutlass.range(self.iterations_pv_k, unroll_full=True): + cute.copy( + tma_atom_scale_v, + tScaleVgV[None, kv_step * self.iterations_pv_k + iter_k], + tScaleVsV[None, (iter_k, scale_v_handle.index)], + tma_bar_ptr=scale_v_handle.barrier, + ) + + for iter_k in cutlass.range(self.iterations_pv_k, unroll_full=True): + for iter_n in cutlass.range(self.iterations_pv_n, unroll_full=True): + v_handle = load_v_producer.acquire_and_advance() + cute.copy( + tma_atom_v, + tVgV[None, iter_n, kv_step * self.iterations_pv_k + iter_k], + tVsV[None, v_handle.index], + tma_bar_ptr=v_handle.barrier, + ) + return load_v_producer, load_scale_v_producer + + @cute.jit + def mma_qk( + self, + 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) + # unroll 4 to avoid spill compared to unroll_full + for iter in cutlass.range(self.iterations_qk, unroll=4): + 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 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 = 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() + for iter_k in cutlass.range(self.iterations_pv_k, unroll=1): + pv_whether_acc = pv_tiled_mma.get(tcgen05.Field.ACCUMULATE) + # unroll 4 to avoid spill compared to unroll_full + for iter_n in cutlass.range(self.iterations_pv_n, unroll=4): + 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_n] + tOrP_slice = tOrP[None, None, None, (iter_k, 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) + 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, sP, smem_exchange = 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) + # warp-wise reduction for lane layout 2x2 + self.smem_exchange_sync_bar.arrive_and_wait() + smem_exchange[thread_idx] = row_max + self.smem_exchange_sync_bar.arrive_and_wait() + row_max = cute.arch.fmax( + row_max, + smem_exchange[ + (thread_idx + 64) % (self.threads_per_warp * len(self.softmax_warp_ids)) + ], + ) + 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 range(0, cute.size(tTMEM_LOADrS), 2): + tTMEM_LOADrS[k], tTMEM_LOADrS[k + 1] = cute.arch.fma_packed_f32x2( + (tTMEM_LOADrS[k], tTMEM_LOADrS[k + 1]), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + ) + tTMEM_LOADrS[k] = cute.math.exp2(tTMEM_LOADrS[k], fastmath=True) + tTMEM_LOADrS[k + 1] = cute.math.exp2(tTMEM_LOADrS[k + 1], 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, (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], sP_slice.shape[3]), + ), + stride=( + (sP_slice.stride[0][0], sP_slice.stride[1]), + (sP_slice.stride[0][1], sP_slice.stride[2], sP_slice.stride[3]), + ), + ), + ) + 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, + value_args: tuple, + o_args: tuple, + epi_tile: cute.Tile, + ) -> pipeline.PipelineConsumer: + (row_max, row_max_prev, scale_softmax_log2) = value_args + (mma_o_consumer, tOtO_staged, cO_staged) = o_args + scale = scale_softmax_log2 * (row_max_prev - row_max) + scale = cute.math.exp2(scale, fastmath=True) + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + o_handle = mma_o_consumer.wait_and_advance() + for iter_n in cutlass.range(self.iterations_pv_n): + tOtO = tOtO_staged[(None, None), 0, 0, iter_n] + cO = cO_staged[None, None, iter_n] + 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(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) + for i in cutlass.range(cute.size(tTMEM_LOADtO, mode=[2]), unroll_full=True): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, 0, i] + tTMEM_STOREtO_i = tTMEM_STOREtO[None, 0, i] + tTMEM_LOADcO_i = tTMEM_LOADcO[None, 0, i] + tTMrO = cute.make_rmem_tensor(tTMEM_LOADcO_i.shape, self.pv_acc_dtype) + cute.copy(tmem_tiled_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), + ) + cute.copy(tmem_store_atom, tTMrO, tTMEM_STOREtO_i) + cute.arch.fence_view_async_tmem_store() + o_handle.release() + return mma_o_consumer + + @cute.jit + def correction_epilog( + self, + value_args: Tuple, + o_args: Tuple, + epi_tile: cute.Tile, + ) -> Tuple[pipeline.PipelineConsumer, pipeline.PipelineProducer]: + (row_sum, seqlen_q, scale_output) = value_args + (mma_o_consumer, gO_staged, cO_staged, tOtO_staged, smem_exchange) = o_args + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + self.smem_exchange_sync_bar.arrive_and_wait() + smem_exchange[thread_idx] = row_sum + self.smem_exchange_sync_bar.arrive_and_wait() + row_sum = ( + row_sum + + smem_exchange[ + (thread_idx + 64) % (self.threads_per_warp * len(self.softmax_warp_ids)) + ] + ) + scale = scale_output / row_sum + o_handle = mma_o_consumer.wait_and_advance() + for iter_n in cutlass.range(self.iterations_pv_n): + 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=[2]), unroll_full=True): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, 0, i] + tTMEM_LOADgO_i = tTMEM_LOADgO[None, 0, i] + tTMEM_LOADcO_i = tTMEM_LOADcO[None, 0, i] + 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 + + @cute.jit + def dequant_k( + self, + tensor_args: Tuple, + pipeline_args: Tuple, + ): + (sOrig, sScale, sTrans) = tensor_args + (load_kv_consumer, load_scale_consumer, dequant_kv_producer) = pipeline_args + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.transform_warp_ids)) + r2s_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.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) + ), + ), + self.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) + ), + ), + self.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 + ) + + # 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(self.q_dtype) + scale = cute.TensorSSA( + tSrScale[None, None, None, None, 0].load(), + transformed_tensor.shape, + self.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, self.iterations_qk, 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], + ) + 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(self.q_dtype) + ) + scale = cute.TensorSSA( + tSrScale[None, None, None, None, iter].load(), + transformed_tensor.shape, + self.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.commit() + + kv_trans_handle = dequant_kv_producer.acquire_and_advance() + cute.autovec_copy( + tTrTrans[None, None, None, None, (self.iterations_qk - 1) % 2], + tTsTrans[None, None, None, None, kv_trans_handle.index], + ) + cute.arch.fence_view_async_shared() + kv_trans_handle.commit() + scale_handle.release() + return load_kv_consumer, load_scale_consumer, dequant_kv_producer + + @cute.jit + def dequant_v( + self, + tensor_args: Tuple, + pipeline_args: Tuple, + mma_tile_coord_v: cute.Coord, + ): + (sOrig, sScale, sTrans) = tensor_args + (load_kv_consumer, load_scale_consumer, dequant_kv_producer) = pipeline_args + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.transform_warp_ids)) + r2s_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.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) + # 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) + ), + ), + self.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) + ), + ), + self.q_dtype, + ) + tSsScale = thr_r2s_tiled_copy.partition_S(sScale) + tSrScale = cute.make_rmem_tensor_like( + tSsScale[None, None, None, None, None, (None, 0)] + ) + scale_v_handle = load_scale_consumer.wait_and_advance() + cute.autovec_copy( + tSsScale[None, None, None, None, None, (None, scale_v_handle.index)], + tSrScale, + ) + # 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(self.q_dtype) + scale = cute.TensorSSA( + tSrScale[None, None, None, None, (mma_tile_coord_v, 0), 0].load(), + transformed_tensor.shape, + self.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, self.iterations_pv, 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], + ) + 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(self.q_dtype) + ) + scale = cute.TensorSSA( + tSrScale[ + None, + None, + None, + None, + (mma_tile_coord_v, iter % self.iterations_pv_n), + iter // self.iterations_pv_n, + ].load(), + transformed_tensor.shape, + self.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.commit() + + kv_trans_handle = dequant_kv_producer.acquire_and_advance() + cute.autovec_copy( + tTrTrans[None, None, None, None, (self.iterations_pv - 1) % 2], + tTsTrans[None, None, None, None, kv_trans_handle.index], + ) + cute.arch.fence_view_async_shared() + kv_trans_handle.commit() + scale_v_handle.release() + + return load_kv_consumer, load_scale_consumer, dequant_kv_producer + + @cute.jit + def get_scale_smem_layout( + self, + 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] * self.d_r,) + tma_view_layout = cute.make_layout( + (mma_tiler[2] * self.d_r), + ) + if cutlass.const_expr(self.scale_granularity < mma_tiler[1]): + # 2cta by default, the rest_mn is at least 1 + rest_mn = size_mn // self.scale_granularity + s2r_view_layout = cute.make_layout( + ( + (self.scale_granularity, rest_mn), + mma_tiler[2], + (2, self.d_r // rest_mn // 2), + ), + stride=((0, 1), self.d_r, (rest_mn, 2 * rest_mn)), + ) + else: + rest_l = self.scale_granularity // mma_tiler[1] + s2r_view_layout = cute.make_layout( + (size_mn, mma_tiler[2], (2, (rest_l, self.d_r))), + stride=(0, self.d_r, (0, (0, 1))), + ) + else: # k + scale_tiler = (mma_tiler[1] * self.d_r,) + tma_view_layout = cute.make_layout((size_mn * self.d_r)) + assert self.scale_granularity % mma_tiler[2] == 0, ( + "scale_granularity must be divisible by mma_tiler[2]" + ) + rest_l = self.scale_granularity // mma_tiler[2] + s2r_view_layout = cute.make_layout( + (size_mn, mma_tiler[2], (rest_l, self.d_r)), + stride=(self.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 + + def scale_tma_partition( + self, + tma_atom_s: cute.CopyAtom, + block_in_cluster_coord: cute.Coord, + scale_cta_layout: cute.Layout, + sS: cute.Tensor, + gS: cute.Tensor, + ) -> tuple[cute.Tensor, cute.Tensor]: + tSsS, tSgS = cute.nvgpu.cpasync.tma_partition( + tma_atom_s, + block_in_cluster_coord, + scale_cta_layout, + sS, + gS, + ) + # Add rest_v mode + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tSsS = cute.make_tensor( + tSsS.iterator, + cute.make_layout( + ((tSsS.layout.shape[0], 1), *tSsS.layout.shape[1:]), + stride=( + (tSsS.layout.stride[0], 0), + *tSsS.layout.stride[1:], + ), + ), + ) + tSgS = cute.make_tensor( + tSgS.iterator, + cute.make_layout( + ((tSgS.layout.shape[0], 1), *tSgS.layout.shape[1:]), + stride=( + (tSgS.layout.stride[0], 0), + *tSgS.layout.stride[1:], + ), + ), + ) + return tSsS, tSgS + + +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], + cta_tiler_mn: Tuple[int, int], + 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 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" cta_tiler_mn: {cta_tiler_mn}") + 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}") + + # 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, 512}: + raise ValueError("head dimension must be 256, or 512") + + if d % scale_granularity != 0: + raise ValueError("head dimension must be divisible by scale_granularity") + + if scale_granularity != d: + raise ValueError("scale_granularity must be equal to head dimension") + + 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) + + cta_tiler_mnk = (*cta_tiler_mn, d) # seq_q, seq_k, d + + mask_type = fmha_utils.MaskEnum.WINDOW_MASK_INFERENCE + if is_causal: + mask_type = fmha_utils.MaskEnum.WINDOW_MASK_INFERENCE + else: + if s_k % cta_tiler_mn[1] != 0: + mask_type = fmha_utils.MaskEnum.RESIDUAL_MASK + + fmha = MixedInputFusedMultiHeadAttentionPrefill( + scale_granularity, + qk_acc_dtype, + pv_acc_dtype, + cta_tiler_mnk, + 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, + ) + + 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 (in bytes)", + ) + + 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( + "--cta_tiler_mn", + type=parse_comma_separated_ints, + default=(64, 256), + help="cta tiler to tile seq_q, seq_k", + ) + + 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 len(args.cta_tiler_mn) != 2: + parser.error("--cta_tiler_mn must contain exactly 2 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.cta_tiler_mn, + 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_gemm.py b/examples/python/CuTeDSL/blackwell/mixed_input_gemm.py index 7d8a4b07..acc88fea 100644 --- a/examples/python/CuTeDSL/blackwell/mixed_input_gemm.py +++ b/examples/python/CuTeDSL/blackwell/mixed_input_gemm.py @@ -27,7 +27,6 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import argparse -from enum import Enum, auto from math import log2, ceil from typing import Optional, Union @@ -37,9 +36,12 @@ 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.torch as cutlass_torch 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 from cutlass.cute.runtime import from_dlpack @@ -130,15 +132,6 @@ Besides the requirements from the Blackwell dense GEMM example, there are some c """ -class TransformMode(Enum): - """ - An enumeration for the possible transform modes of a mixed-input GEMM. - """ - - ConvertOnly = auto() - ConvertScale = auto() - - class MixedInputGemmKernel: """ Mixed-input GEMM kernel for NVIDIA Blackwell SM100 architecture. @@ -226,7 +219,7 @@ class MixedInputGemmKernel: + 1 ) - # Set barrier id for cta sync, epilogue sync, tmem ptr sync, and transform sync + # 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) ) @@ -234,7 +227,6 @@ class MixedInputGemmKernel: 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.smem_buffer_align_bytes = 1024 @@ -255,7 +247,9 @@ class MixedInputGemmKernel: - Computing tensor memory allocation columns """ # Deduce where the transformed A tensor is stored, shared memory(SMEM) or tensor memory(TMEM) - self.transform_a_source = self._get_transform_a_source(self.a_major_mode) + 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, @@ -346,7 +340,7 @@ class MixedInputGemmKernel: self.smem_layout_a, self.smem_layout_a_transform, self.smem_layout_b, - ) = self._compute_smem_layout( + ) = mixed_input_utils.compute_smem_layout( tiled_mma, self.mma_tiler, self.a_dtype, @@ -358,11 +352,20 @@ class MixedInputGemmKernel: self.smem_layout_scale_per_stage = None self.smem_layout_scale = None if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): - # Get smem layout for scale tensor + # Get scale tile shape and smem layout for scale tensor ( + self.scale_tile_shape, self.smem_layout_scale_per_stage, self.smem_layout_scale, - ) = self.get_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, @@ -448,7 +451,12 @@ class MixedInputGemmKernel: 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 = self.get_gmem_layout_scale(a.shape) + 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) @@ -466,8 +474,12 @@ class MixedInputGemmKernel: self.transform_a_source, ) # Set up gmem copy atoms for A, scale, and B - a_op = self._get_tma_atom_kind(self.is_a_mcast, self.use_2cta_instrs, False) - b_op = self._get_tma_atom_kind(self.is_b_mcast, self.use_2cta_instrs, True) + 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)) @@ -650,7 +662,6 @@ class MixedInputGemmKernel: grid=grid, block=[self.threads_per_cta, 1, 1], cluster=(*self.cluster_shape_mn, 1), - smem=self.shared_storage.size_in_bytes(), stream=stream, min_blocks_per_mp=1, ) @@ -730,6 +741,7 @@ class MixedInputGemmKernel: 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 @@ -753,6 +765,7 @@ class MixedInputGemmKernel: 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 @@ -766,6 +779,7 @@ class MixedInputGemmKernel: ), 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 @@ -779,6 +793,7 @@ class MixedInputGemmKernel: 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 @@ -790,6 +805,7 @@ class MixedInputGemmKernel: 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 @@ -802,8 +818,7 @@ class MixedInputGemmKernel: ) # Cluster arrive after barrier init - if cutlass.const_expr(cute.size(self.cluster_shape_mn) > 1): - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # Setup smem tensor A/scale/B/C sC = ( @@ -897,7 +912,7 @@ class MixedInputGemmKernel: 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 = self._get_copy_atom_a_transform( + copy_atom_a_transform = mixed_input_utils.get_copy_atom_a_transform( self.mma_dtype, self.use_2cta_instrs, self.transform_a_source, @@ -928,7 +943,7 @@ class MixedInputGemmKernel: tCsS = thr_mma.partition_A(sS_input) # ((atom_v, rest_v), STAGE) # ((atom_v, rest_v), loopM, loopK, loopL) - tSsS, tSgS = self.scale_tma_partition( + tSsS, tSgS = mixed_input_utils.scale_tma_partition( tCsS, tCgS, tma_atom_s, @@ -959,10 +974,7 @@ class MixedInputGemmKernel: ) # Cluster wait before TMEM alloc and ensure pipelines are ready - if cutlass.const_expr(cute.size(self.cluster_shape_mn) > 1): - cute.arch.cluster_wait() - else: - self.cta_sync_barrier.arrive_and_wait() + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) # TMEM allocation tmem.allocate(self.num_tmem_alloc_cols) @@ -1145,7 +1157,7 @@ class MixedInputGemmKernel: dst_copy_a, tAsA_input, tAsA_transform, - ) = self.transform_partition( + ) = mixed_input_utils.transform_partition( self.transform_a_source, self.scale_mode, copy_atom_a_input, @@ -1173,8 +1185,10 @@ class MixedInputGemmKernel: tSrS_copy = None tSrS = None if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): - smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS = self.scale_partition( - src_copy_a, tCsS, transform_local_tidx, self.mma_dtype + 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" @@ -1582,328 +1596,6 @@ class MixedInputGemmKernel: if cutlass.const_expr(self.use_tma_store): c_pipeline.producer_tail() - def scale_tma_partition( - self, - tCsS: cute.Tensor, - tCgS: cute.Tensor, - tma_atom_s: cute.CopyAtom, - block_in_cluster_coord_vmnk: cute.Coord, - scale_cta_layout: cute.Layout, - ) -> tuple[cute.Tensor, cute.Tensor]: - """ - Perform TMA partition for scale tensor. - - This method partitions the gobal memory and shared memory buffer for scale tensor for TMA load. - - :param tCsS: Input scale shared memory tensor - :type tCsS: cute.Tensor - :param tCgS: Input scale global memory tensor - :type tCgS: cute.Tensor - :param tma_atom_s: TMA copy atom for scale tensor - :type tma_atom_s: cute.CopyAtom - :param block_in_cluster_coord_vmnk: CTA coord in the cluster - :type block_in_cluster_coord_vmnk: cute.Coord - :param scale_cta_layout: Layout of CTA from the view of the scale tensor - :type scale_cta_layout: cute.Layout - - :return: A tuple containing (tSsS, tSgS) where: - - tSsS: Partitioned scale tensor in shared memory - - tSgS: Partitioned scale tensor in global memory - :rtype: tuple[cute.Tensor, cute.Tensor] - """ - tSsS, tSgS = cpasync.tma_partition( - tma_atom_s, - block_in_cluster_coord_vmnk[2], - scale_cta_layout, - cute.group_modes(tCsS, 0, 3), - cute.group_modes(tCgS, 0, 3), - ) - # Add rest_v mode - # ((atom_v, rest_v), STAGE) - # ((atom_v, rest_v), loopM, loopK, loopL) - tSsS = cute.make_tensor( - tSsS.iterator, - cute.make_layout( - ((tSsS.layout.shape[0], 1), *tSsS.layout.shape[1:]), - stride=( - (tSsS.layout.stride[0], 0), - *tSsS.layout.stride[1:], - ), - ), - ) - tSgS = cute.make_tensor( - tSgS.iterator, - cute.make_layout( - ((tSgS.layout.shape[0], 1), *tSgS.layout.shape[1:]), - stride=( - (tSgS.layout.stride[0], 0), - *tSgS.layout.stride[1:], - ), - ), - ) - return tSsS, tSgS - - def transform_partition( - self, - transform_a_source: tcgen05.OperandSource, - scale_mode: TransformMode, - copy_atom_a_input: cute.CopyAtom, - copy_atom_a_transform: cute.CopyAtom, - sA_input: cute.Tensor, - A_transform: cute.Tensor, - transform_local_tidx: cutlass.Int32, - ) -> tuple[cute.TiledCopy, cute.TiledCopy, cute.Tensor, cute.Tensor]: - """ - Partition tensors for transform input and output. - - This method sets up the copy atoms and partitions the shared/tensor memory - for the transformation of tensor A. - - :param transform_a_source: Where the transformed tensor A is stored (TMEM or SMEM) - :type transform_a_source: tcgen05.OperandSource - :param scale_mode: The transform mode (ConvertOnly or ConvertScale) - :type scale_mode: TransformMode - :param copy_atom_a_input: Copy atom for loading A from shared memory - :type copy_atom_a_input: cute.CopyAtom - :param copy_atom_a_transform: Copy atom for storing transformed A - :type copy_atom_a_transform: cute.CopyAtom - :param sA_input: Input tensor A in shared memory - :type sA_input: cute.Tensor - :param A_transform: Transformed tensor A in tensor or shared memory - :type A_transform: cute.Tensor - :param transform_local_tidx: Local thread index for transformation warps - :type transform_local_tidx: cutlass.Int32 - - :return: A tuple containing (src_copy_a, dst_copy_a, tAsA_input, tA_transform) where: - - src_copy_a: Tiled copy for source tensor - - dst_copy_a: Tiled copy for destination tensor - - tAsA_input: Partitioned input tensor A - - tA_transform: Partitioned transformed tensor A - :rtype: tuple[cute.TiledCopy, cute.TiledCopy, cute.Tensor, cute.Tensor] - """ - if cutlass.const_expr(transform_a_source == tcgen05.OperandSource.TMEM): - if cutlass.const_expr( - cute.size(A_transform, mode=[0, 0]) == 128 - and cute.size(sA_input, mode=[0, 0]) == 64 - ): - tensor_input = cute.make_tensor( - sA_input.iterator, - cute.logical_product( - sA_input.layout, - ((cute.make_layout(2, stride=0), None), None, None, None), - ), - ) - else: - tensor_input = sA_input - reg2tmem_tiled_copy = tcgen05.make_tmem_copy( - copy_atom_a_transform, A_transform[(None, None, None, 0)] - ) - thr_reg2tmem_tiled_copy = reg2tmem_tiled_copy.get_slice( - transform_local_tidx - ) - partitioned_tensor_input = thr_reg2tmem_tiled_copy.partition_S(tensor_input) - partitioned_tensor_transform = thr_reg2tmem_tiled_copy.partition_D( - A_transform - ) - src_copy_a = ( - cute.make_tiled_copy_S(copy_atom_a_input, reg2tmem_tiled_copy) - if scale_mode is TransformMode.ConvertScale - else None - ) - dst_copy_a = reg2tmem_tiled_copy - tAsA_input = partitioned_tensor_input - tA_transform = partitioned_tensor_transform - elif cutlass.const_expr(transform_a_source == tcgen05.OperandSource.SMEM): - # Construct tiled_copy satisfying 8 contiguous elts per copy atom - reg2smem_tiled_copy = cute.make_cotiled_copy( - copy_atom_a_transform, - cute.make_layout((128, 8), stride=(8, 1)), - A_transform[(None, None, None, 0)].layout, - ) - thr_reg2smem_tiled_copy = reg2smem_tiled_copy.get_slice( - transform_local_tidx - ) - partitioned_tensor_input = thr_reg2smem_tiled_copy.partition_S(sA_input) - partitioned_tensor_transform = thr_reg2smem_tiled_copy.partition_D( - A_transform - ) - src_copy_a = ( - cute.make_tiled_copy_S(copy_atom_a_input, reg2smem_tiled_copy) - if scale_mode is TransformMode.ConvertScale - else None - ) - # auto-vec copy is enough for copy from register to shared memory here - dst_copy_a = None - tAsA_input = partitioned_tensor_input - tA_transform = partitioned_tensor_transform - return src_copy_a, dst_copy_a, tAsA_input, tA_transform - - def scale_partition( - self, - src_copy_a: cute.TiledCopy, - tCsS: cute.Tensor, - transform_local_tidx: cutlass.Int32, - mma_dtype: type[cutlass.Numeric], - ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor]: - """ - Partition the scale tensor for transformation. - - This method prepares the copy atom and partitions the shared memory for the scale tensor. - - :param src_copy_a: Tiled copy for the source tensor - :type src_copy_a: cute.TiledCopy - :param tCsS: Scale tensor in shared memory - :type tCsS: cute.Tensor - :param transform_local_tidx: Local thread index for transformation warps - :type transform_local_tidx: cutlass.Int32 - :param mma_dtype: Data type for the MMA operation - :type mma_dtype: type[cutlass.Numeric] - - :return: A tuple containing (smem_thr_copy_S, tSsS_trans, tSrS) where: - - smem_thr_copy_S: Tiled copy for the scale tensor - - tSsS_trans: Partitioned scale tensor for transformation - - tSrS_copy: Register fragment for the scale tensor - - tSrS: view of scale tensor used for transformation computation - :rtype: tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor] - """ - smem_thr_copy_S = None - tSsS_trans = None - tSrS = None - # Partition scale tensor - smem_thr_copy_S = src_copy_a.get_slice(transform_local_tidx) - tSsS_trans = smem_thr_copy_S.partition_S(tCsS) - # Construct register fragment for scale tensor - tSsS_layout_per_stage = tSsS_trans[(None, None, None, None, 0)].layout - # tSrS for copy - tSrS_copy = cute.make_rmem_tensor( - cute.filter_zeros(tSsS_layout_per_stage).shape, mma_dtype - ) - # tSrS view for transformation computation - tSrS = cute.make_tensor( - tSrS_copy.iterator, - cute.make_layout( - tSsS_layout_per_stage.shape, stride=tSrS_copy.layout.stride - ), - ) - return smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS - - def get_gmem_layout_scale( - self, scale_shape_mkl: tuple[int, int, int] - ) -> cute.Layout: - """ - Get the layout of the scale tensor in global memory. - - :param scale_shape_mkl: The shape of the scale tensor (M, K, L). - :type scale_shape_mkl: tuple[int, int, int] - - :return: The layout of the scale tensor in global memory. - :rtype: cute.Layout - """ - m, k, l = scale_shape_mkl - shape_scale = ( - (self.scale_granularity_m, cute.ceil_div(m, self.scale_granularity_m)), - (self.scale_granularity_k, cute.ceil_div(k, self.scale_granularity_k)), - ) - if cutlass.const_expr(self.scale_major_mode == tcgen05.OperandMajorMode.MN): - layout_mk = cute.make_layout( - shape_scale, - stride=( - (0, 1), - (0, cute.size(shape_scale[0][1])), - ), - ) - else: - layout_mk = cute.make_layout( - shape_scale, - stride=( - (0, cute.size(shape_scale[1][1])), - (0, 1), - ), - ) - return cute.make_layout( - (*layout_mk.shape, l), - stride=(*layout_mk.stride, cute.cosize(layout_mk)), - ) - - def get_smem_layout_scale(self) -> tuple[cute.ComposedLayout, cute.ComposedLayout]: - """ - Get the layout of the scale tensor in shared memory. - - :return: A tuple containing: - - smem_layout_scale_per_stage: Shared memory layout for scale tensor per stage - - smem_layout_scale: Shared memory layout for scale tensor - :rtype: tuple[cute.ComposedLayout, cute.ComposedLayout] - """ - self.scale_tile_shape = ( - ( - cute.size(self.mma_tiler[0]) // 2 - if self.use_2cta_instrs - else cute.size(self.mma_tiler[0]) - ), - cute.size(self.mma_tiler[2]), - ) - size_mn = self.scale_tile_shape[0] - size_k = self.scale_tile_shape[1] - smem_size_mn = ( - self.scale_granularity_m if self.scale_granularity_m < size_mn else size_mn - ) - smem_size_k = ( - self.scale_granularity_k if self.scale_granularity_k < size_k else size_k - ) - div_mn = cute.ceil_div(size_mn, smem_size_mn) - div_k = cute.ceil_div(size_k, smem_size_k) - smem_atom_shape = ( - (smem_size_mn, div_mn), - (smem_size_k, div_k), - ) - if cutlass.const_expr(self.scale_major_mode == tcgen05.OperandMajorMode.MN): - outer_layout = cute.make_layout( - smem_atom_shape, - stride=( - (0, 1), - (0, div_mn), - ), - ) - else: - outer_layout = cute.make_layout( - smem_atom_shape, - stride=( - (0, div_k), - (0, 1), - ), - ) - # Apply a trivial swizzle to make it a composed layout, which could be used to construct TMA atom - smem_layout_scale_per_stage = cute.make_composed_layout( - cute.make_swizzle(0, 4, 3), 0, outer_layout - ) - assert cute.rank(smem_layout_scale_per_stage) == 2, ( - "Scale layout must be rank 2" - ) - - assert ( - cute.size(self.mma_tiler[0]) - % cute.size(smem_layout_scale_per_stage.outer[0]) - == 0 - ), "smem_layout_scale_per_stage must equal the tile shape." - assert ( - cute.size(self.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(self.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, - cute.make_layout( - (self.num_scale_load2trans_stage), - stride=(cute.cosize(smem_layout_scale_per_stage.outer)), - ), - ) - return smem_layout_scale_per_stage, smem_layout_scale - def epilog_gmem_copy_and_partition( self, tidx: cutlass.Int32, @@ -2285,126 +1977,6 @@ class MixedInputGemmKernel: num_tmem_a_cols, ) - @staticmethod - def _compute_smem_layout( - tiled_mma: cute.TiledMma, - mma_tiler_mnk: tuple[int, int, int], - a_dtype: type[cutlass.Numeric], - b_dtype: type[cutlass.Numeric], - load2trans_stage_count: int, - trans2mma_stage_count: int, - ) -> tuple[ - cute.ComposedLayout, - cute.ComposedLayout, - cute.ComposedLayout, - ]: - """ - Compute shared memory layouts for tensor A, transformed A and tensor B. - - :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 load2trans_stage_count: Number of stages for load-to-transform pipeline. - :type load2trans_stage_count: int - :param trans2mma_stage_count: Number of stages for transform-to-MMA pipeline. - :type trans2mma_stage_count: int - - :return: A tuple containing: - - smem_layout_a: Shared memory layout for tensor A - - smem_layout_a_transform: Shared memory layout for transformed tensor A - - smem_layout_b: Shared memory layout for tensor B - :rtype: tuple[cute.ComposedLayout, cute.ComposedLayout, cute.ComposedLayout] - """ - smem_layout_a = sm100_utils.make_smem_layout_a( - tiled_mma, - mma_tiler_mnk, - a_dtype, - load2trans_stage_count, - ) - smem_layout_a_transform = sm100_utils.make_smem_layout_a( - tiled_mma, - mma_tiler_mnk, - tiled_mma.op.a_dtype, - trans2mma_stage_count, - ) - smem_layout_b = sm100_utils.make_smem_layout_b( - tiled_mma, - mma_tiler_mnk, - b_dtype, - load2trans_stage_count, - ) - return ( - smem_layout_a, - smem_layout_a_transform, - smem_layout_b, - ) - - @staticmethod - def _get_transform_a_source( - a_major_mode: tcgen05.OperandMajorMode, - ) -> tcgen05.OperandSource: - """ - Determine the operand source for transformed A tensor based on the operand major mode. - """ - if cutlass.const_expr(a_major_mode == tcgen05.OperandMajorMode.K): - return tcgen05.OperandSource.TMEM - else: - return tcgen05.OperandSource.SMEM - - @staticmethod - def _get_tma_atom_kind( - mcast: cutlass.Boolean, - use_2cta_instrs: bool, - is_b: bool, - ) -> Union[ - cpasync.CopyBulkTensorTileG2SMulticastOp, cpasync.CopyBulkTensorTileG2SOp - ]: - """ - Get the TMA atom kind based on 1) whether it's a multicast operation, - 2) whether 2CTA tcgen05.mma instruction is enabled, and - 3) whether it's a B tensor - """ - # Not using .2CTA instructions for tensor A as the consumer is threads on different CTAs - cta_group = ( - tcgen05.CtaGroup.TWO if (use_2cta_instrs and is_b) else tcgen05.CtaGroup.ONE - ) - if cutlass.const_expr(mcast): - return cpasync.CopyBulkTensorTileG2SMulticastOp(cta_group) - return cpasync.CopyBulkTensorTileG2SOp(cta_group) - - @staticmethod - def _get_copy_atom_a_transform( - mma_dtype: type[cutlass.Numeric], - use_2cta_instrs: bool, - transform_a_source: tcgen05.OperandSource, - a_smem_shape: cute.Shape, - a_dtype: type[cutlass.Numeric], - ) -> cute.CopyAtom: - """ - Determine the copy atom for transformed A tensor based on the operand source and tile size. - """ - if cutlass.const_expr(transform_a_source == tcgen05.OperandSource.TMEM): - if cutlass.const_expr( - cute.size(a_smem_shape[0][0]) == 64 and (not use_2cta_instrs) - ): - copy_op_r2t = tcgen05.St16x256bOp( - tcgen05.Repetition(1), tcgen05.Unpack.NONE - ) - else: - copy_op_r2t = tcgen05.St32x32bOp( - tcgen05.Repetition(8), tcgen05.Unpack.NONE - ) - return cute.make_copy_atom(copy_op_r2t, mma_dtype) - else: - return cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), a_dtype, num_bits_per_copy=32 - ) - @staticmethod def _compute_grid( c: cute.Tensor, @@ -2429,29 +2001,6 @@ class MixedInputGemmKernel: return tile_sched_params, grid - def is_valid_scale_granularity( - scale_granularity_m: int, - scale_granularity_k: int, - a_dtype: type[cutlass.Numeric], - k: int, - mma_tiler_k: int, - ) -> bool: - """ - Check if the scale granularity settings are valid for the given data type and problem size. - """ - if a_dtype.width == 8: - # No scale tensor for 8bit data type A - if not (scale_granularity_m == 0 and scale_granularity_k == 0): - return False - elif a_dtype.width == 4: - if scale_granularity_m != 1 or ( - scale_granularity_k == 0 - or k % scale_granularity_k != 0 - or scale_granularity_k % mma_tiler_k != 0 - ): - return False - return True - def is_valid_tensor_alignment( m: int, n: int, @@ -2566,7 +2115,7 @@ class MixedInputGemmKernel: mma_tiler, cluster_shape_mn, use_2cta_instrs ): return False - if not MixedInputGemmKernel.is_valid_scale_granularity( + if not mixed_input_utils.is_valid_scale_granularity( scale_granularity_m, scale_granularity_k, a_dtype, k, mma_tiler[2] ): return False @@ -2634,7 +2183,11 @@ def create_i4_tensor_and_scale( 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) + a_max = ( + torch.maximum(ref / up_4b, ref / lb_4b) + if dtype == cutlass.Int4 + else torch.maximum(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 @@ -2668,17 +2221,6 @@ def create_i4_tensor_and_scale( ) -def get_divisibility(contiguous_dim_size: int, upper_bound: int = 128) -> int: - """ - Calculate the largest power of 2 divisibility factor for memory alignment. - """ - # Check the largest power of 2 factor of contiguous_dim_size - for i in range(int(log2(contiguous_dim_size)), 0, -1): - if contiguous_dim_size % (2**i) == 0: - return min(2**i, upper_bound) - return 1 - - def create_tensor_a( l: int, m: int, @@ -2688,7 +2230,7 @@ def create_tensor_a( scale_granularity_m: int = 0, scale_granularity_k: int = 0, transformed_dtype: Optional[type[cutlass.Numeric]] = None, -) -> tuple[cute.Tensor, cute.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[cute.Tensor, Optional[cute.Tensor], torch.Tensor, Optional[torch.Tensor]]: """ Create tensor A and scale tensor. """ @@ -2710,7 +2252,7 @@ def create_tensor_a( a_dtype, scale_granularity_m, scale_granularity_k, - divisibility=get_divisibility(m if a_major == "m" else k), + divisibility=mixed_input_utils.get_divisibility(m if a_major == "m" else k), transformed_dtype=transformed_dtype, ) else: @@ -2725,7 +2267,9 @@ def create_tensor_a( a_torch_cpu, a_dtype, is_dynamic_layout=True, - assumed_align=get_divisibility(m if a_major == "m" else k), + 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 @@ -2774,18 +2318,18 @@ def create_tensors( b_torch_cpu, b_dtype, is_dynamic_layout=True, - assumed_align=get_divisibility(n if b_major == "n" else k), + 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=get_divisibility(m if c_major == "m" else n), + 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=get_divisibility(m if c_major == "m" else n), + divisibility=mixed_input_utils.get_divisibility(m if c_major == "m" else n), ) return ( @@ -2966,25 +2510,36 @@ def run( def generate_tensors(): a_tensor, a_scale_tensor, a_torch_cpu, a_scale_torch_cpu = create_tensor_a( - l, m, k, a_major, a_dtype, scale_granularity_m, scale_granularity_k, b_dtype + l, + m, + k, + a_major, + a_dtype, + scale_granularity_m, + scale_granularity_k, + b_dtype, ) b_tensor, _ = cutlass_torch.cute_tensor_like( b_torch_cpu, b_dtype, is_dynamic_layout=True, - assumed_align=get_divisibility(n if b_major == "n" else k), + assumed_align=mixed_input_utils.get_divisibility( + n if b_major == "n" else k + ), ) c_torch_cpu = cutlass_torch.matrix(l, m, 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=get_divisibility(m if c_major == "m" else n), + 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=get_divisibility(m if c_major == "m" else n), + divisibility=mixed_input_utils.get_divisibility(m if c_major == "m" else n), ) return testing.JitArguments( a_tensor, a_scale_tensor, b_tensor, c_tensor, current_stream diff --git a/examples/python/CuTeDSL/blackwell/mla.py b/examples/python/CuTeDSL/blackwell/mla.py index 4bdbcac5..7f259c6d 100644 --- a/examples/python/CuTeDSL/blackwell/mla.py +++ b/examples/python/CuTeDSL/blackwell/mla.py @@ -43,6 +43,7 @@ import cutlass.cute.nvgpu.tcgen05 as tcgen05 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.torch as cutlass_torch import cutlass.utils.blackwell_helpers as sm100_utils from cutlass.cute.runtime import from_dlpack @@ -970,7 +971,6 @@ class BlackwellMultiHeadLatentAttentionForward: num_tmem_dealloc_threads = self.threads_per_warp * self.num_compute_warps with cute.arch.elect_one(): cute.arch.mbarrier_init(tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads) - cute.arch.mbarrier_init_fence() load_q_pipeline = self.make_and_init_load_qkv_pipeline( storage.load_q_mbar_ptr.data_ptr(), @@ -1004,8 +1004,7 @@ class BlackwellMultiHeadLatentAttentionForward: ) # Cluster arrive after barrier init - if cutlass.const_expr(cute.size(self.cluster_shape_mnk) > 1): - cute.arch.cluster_arrive_relaxed() + 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) @@ -1035,10 +1034,7 @@ class BlackwellMultiHeadLatentAttentionForward: # # Cluster wait before tensor memory alloc # - if cutlass.const_expr(cute.size(self.cluster_shape_mnk) > 1): - cute.arch.cluster_wait() - else: - pipeline.sync(barrier_id=4) + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk) # /////////////////////////////////////////////////////////////////////////////// # Load warps, including page table and data tensors @@ -2046,8 +2042,9 @@ class BlackwellMultiHeadLatentAttentionForward: # wait cpasync arrive until the last stage load_kv_pipeline = common_params.load_kv_pipeline - if copy_in_flight_count == self.load_kv_stage: - cute.arch.cp_async_wait_group(self.load_kv_stage - 1) + release_distance = 2 + if copy_in_flight_count == self.load_kv_stage - release_distance: + cute.arch.cp_async_wait_group(self.load_kv_stage - release_distance - 1) load_kv_pipeline.producer_commit(load_kv_commit_state) load_kv_commit_state.advance() copy_in_flight_count -= 1 @@ -3922,6 +3919,7 @@ class BlackwellMultiHeadLatentAttentionForward: producer_group=load_qkv_producer_group, consumer_group=load_qkv_consumer_group, cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, ) else: load_qkv_producer_group = pipeline.CooperativeGroup( @@ -3937,6 +3935,7 @@ class BlackwellMultiHeadLatentAttentionForward: 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( @@ -3971,6 +3970,7 @@ class BlackwellMultiHeadLatentAttentionForward: 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( @@ -4005,6 +4005,7 @@ class BlackwellMultiHeadLatentAttentionForward: 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( @@ -4033,6 +4034,7 @@ class BlackwellMultiHeadLatentAttentionForward: 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( @@ -4067,6 +4069,7 @@ class BlackwellMultiHeadLatentAttentionForward: producer_group=mma_o_producer_group, consumer_group=mma_o_consumer_group, cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, ) def make_and_init_load_pt_pipeline(self, load_pt_mbar_ptr): @@ -4091,6 +4094,7 @@ class BlackwellMultiHeadLatentAttentionForward: num_stages=self.load_pt_stage, producer_group=load_pt_producer_group, consumer_group=load_pt_consumer_group, + defer_sync=True, ) @staticmethod diff --git a/examples/python/CuTeDSL/cute/fake_tensor.py b/examples/python/CuTeDSL/cute/fake_tensor.py new file mode 100644 index 00000000..fa53c506 --- /dev/null +++ b/examples/python/CuTeDSL/cute/fake_tensor.py @@ -0,0 +1,48 @@ +import cutlass +import cutlass.cute as cute + + +""" +Example of using fake tensors in CuTe. + +This script demonstrates how to use fake tensors in CuTe to drive compilation without creating actual tensors +from frameworks like PyTorch or TensorFlow. + +Run this file directly to see the output type information. +""" + + +@cute.jit +def print_tensor_type(t: cute.Tensor): + print(t) + + +def run(): + from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_tensor + + shape = (3, 4) + a = make_fake_compact_tensor(cutlass.Float16, (3, 4), stride_order=(1, 0)) + cute.compile(print_tensor_type, a) + + # 32-bit symbolic integer with divisibility 8 + shape = (3, cute.sym_int32(divisibility=8)) + a = make_fake_compact_tensor(cutlass.Float16, shape, stride_order=(1, 0)) + cute.compile(print_tensor_type, a) + + # with static stride + a = make_fake_tensor(cutlass.Float16, shape, stride=(4, 1)) + cute.compile(print_tensor_type, a) + + # with dynamic stride using 32bit integer + stride = (cute.sym_int32(divisibility=8), 1) + a = make_fake_tensor(cutlass.Float16, shape, stride=stride) + cute.compile(print_tensor_type, a) + + # with dynamic stride using 64bit integer + stride = (cute.sym_int64(divisibility=8), 1) + a = make_fake_tensor(cutlass.Float16, shape, stride=stride) + cute.compile(print_tensor_type, a) + + +if __name__ == "__main__": + run() diff --git a/examples/python/CuTeDSL/cute/ffi/CMakeLists.txt b/examples/python/CuTeDSL/cute/ffi/CMakeLists.txt index 1870197f..1d40b32f 100644 --- a/examples/python/CuTeDSL/cute/ffi/CMakeLists.txt +++ b/examples/python/CuTeDSL/cute/ffi/CMakeLists.txt @@ -35,20 +35,26 @@ find_package(Python3 COMPONENTS Interpreter Development REQUIRED) # Get Python site-packages directory using Python execute_process( - COMMAND ${Python3_EXECUTABLE} -c "import site; print(site.getsitepackages()[0])" - OUTPUT_VARIABLE Python_SITE_PACKAGES + COMMAND ${Python3_EXECUTABLE} -c "import sys, sysconfig; print(';'.join([sysconfig.get_paths()['purelib'],sysconfig.get_paths(vars={'base': sys.base_prefix, 'platbase': sys.base_prefix})['purelib']]))" + OUTPUT_VARIABLE Python_SITE_PACKAGES_PATHS OUTPUT_STRIP_TRAILING_WHITESPACE ) -message(STATUS "Python site-packages directory: ${Python_SITE_PACKAGES}") +message(STATUS "Python site-packages directories: ${Python_SITE_PACKAGES_PATHS}") # Add nanobind path to CMAKE_PREFIX_PATH -list(APPEND CMAKE_PREFIX_PATH ${Python_SITE_PACKAGES}/nanobind/cmake) +foreach(path IN LISTS Python_SITE_PACKAGES_PATHS) + if(EXISTS "${path}/nanobind/cmake") + message(STATUS "Adding nanobind cmake path: ${path}/nanobind/cmake") + list(APPEND CMAKE_PREFIX_PATH "${path}/nanobind/cmake") + break() + endif() +endforeach() # Find nanobind find_package(nanobind) if(NOT nanobind_FOUND) - message(FATAL_ERROR + message(FATAL_ERROR "nanobind not found!\n" "Please install nanobind with: pip install nanobind\n" ) diff --git a/examples/python/CuTeDSL/cute/ffi/jit_argument.py b/examples/python/CuTeDSL/cute/ffi/jit_argument.py index bf21bc17..1a7d6b88 100644 --- a/examples/python/CuTeDSL/cute/ffi/jit_argument.py +++ b/examples/python/CuTeDSL/cute/ffi/jit_argument.py @@ -243,12 +243,7 @@ import tempfile import torch -def run_test(tmpdir=None, cmake_args=""): - # Skip cleanup if user provides tmpdir - cleanup = tmpdir is None - # Initialize temporary build directory - tmpdir = tmpdir or tempfile.mkdtemp() - +def run_test(tmpdir=None, cmake_args="", cleanup=True): try: current_dir = os.path.dirname(os.path.abspath(__file__)) @@ -256,8 +251,6 @@ def run_test(tmpdir=None, cmake_args=""): subprocess.run(["cmake", "-B", tmpdir, current_dir] + cmake_args, check=True) subprocess.run(["cmake", "--build", tmpdir], check=True) - sys.path.append(tmpdir) - from tensor import make_tensor, pycapsule_get_pointer # Mock test tensor and corresponding C structure for this example @@ -314,4 +307,13 @@ if __name__ == "__main__": ) args = parser.parse_args() - run_test(tmpdir=args.tmp_dir, cmake_args=args.cmake_args) + if args.tmp_dir: + tmp_dir = args.tmp_dir + cleanup = False + else: + tmp_dir = tempfile.mkdtemp() + cleanup = True + + sys.path.append(tmp_dir) + + run_test(tmp_dir, args.cmake_args, cleanup) diff --git a/examples/python/CuTeDSL/cute/print_latex.py b/examples/python/CuTeDSL/cute/print_latex.py new file mode 100644 index 00000000..547effeb --- /dev/null +++ b/examples/python/CuTeDSL/cute/print_latex.py @@ -0,0 +1,72 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import argparse +import cutlass +import cutlass.cute as cute +from cutlass.utils import print_latex, print_latex_tv +from cutlass import for_generate, yield_out + +""" +A Latex Printing Example using CuTe DSL. +This example prints latex for a given layout or thread value layout. +The primary goal for this example is to demonstrate how to dump latex, which can then be +turned into an image in your favorite latex compiler. +To run this example: +.. code-block:: bash + python examples/python/CuteDSL/cute/print_latex.py + python examples/python/CuteDSL/cute/print_latex.py --tv_layout +To compile, pipe the output to a file and use a tool like pdflatex: +.. code-block:: bash + python examples/python/CuTeDSL/cute/print_latex.py > latex.tex + pdflatex latex.tex +""" + + +@cute.jit +def main(print_tv_layout: cutlass.Constexpr[bool]): + # Note: only support compile time printing layouts + if cutlass.const_expr(print_tv_layout): + thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0)) + val_layout = cute.make_ordered_layout((4, 1), order=(1, 0)) + tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout) + print_latex_tv(tv_layout, tiler_mn) + else: + layout = cute.make_layout((10, 10)) + print_latex(layout) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="example of print latex and print latex tv" + ) + parser.add_argument("--tv_layout", action="store_true") + + args = parser.parse_args() + + main(args.tv_layout) diff --git a/examples/python/CuTeDSL/cute/torch_fake_tensor.py b/examples/python/CuTeDSL/cute/torch_fake_tensor.py index 06c51ea8..0437bb3d 100644 --- a/examples/python/CuTeDSL/cute/torch_fake_tensor.py +++ b/examples/python/CuTeDSL/cute/torch_fake_tensor.py @@ -65,7 +65,7 @@ def print_tensor(t: cute.Tensor): cute.print_tensor(t) -if __name__ == "__main__": +def run(): from torch._subclasses.fake_tensor import FakeTensorMode shape = (3, 4) @@ -75,3 +75,7 @@ if __name__ == "__main__": real_tensor = torch.randn(shape, dtype=torch.float32) compiled_fn(from_dlpack(real_tensor)) + + +if __name__ == "__main__": + run() diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py b/examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py new file mode 100644 index 00000000..b1d7a97c --- /dev/null +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py @@ -0,0 +1,92 @@ +# 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. + +"""Example demonstrating how to use TVM-FFI ABI with CuTe. + +This example shows how to: +1. Compile a CuTe function with "--enable-tvm-ffi" option +2. Export the compiled function to a shared library +3. Load the shared library and use the compiled function to work with torch.Tensor + +To run this example: + +.. code-block:: bash + + python examples/cute/tvm_ffi/aot_export.py + # run example to use in torch + python examples/cute/tvm_ffi/aot_use_in_torch.py + # run example to use in jax + python examples/cute/tvm_ffi/aot_use_in_jax.py + # run example to use in c++ bundle + bash examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh +""" + +from pathlib import Path +import torch +import os +import subprocess +import tvm_ffi +import torch +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack + + +@cute.kernel +def device_add_one(a: cute.Tensor, b: cute.Tensor): + for i in range(a.shape[0]): + b[i] = a[i] + 1 + + +@cute.jit +def add_one(a: cute.Tensor, b: cute.Tensor): + """b = a + 1""" + device_add_one(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1)) + + +def main(): + # compile the kernel with "--enable-tvm-ffi" option + a_torch = torch.arange(10, dtype=torch.float32, device="cuda") + b_torch = torch.zeros(10, dtype=torch.float32, device="cuda") + a_cute = from_dlpack(a_torch, enable_tvm_ffi=True).mark_layout_dynamic() + b_cute = from_dlpack(b_torch, enable_tvm_ffi=True).mark_layout_dynamic() + # compile the kernel with "--enable-tvm-ffi" option + compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi") + + object_file_path = "./build/add_one.o" + lib_path = "./build/add_one.so" + compiled_add_one.export_to_c(object_file_path, function_name="add_one") + shared_libs = cute.runtime.find_runtime_libraries(enable_tvm_ffi=True) + # compile the object file to a shared library + cmd = ["gcc", "-shared", "-o", lib_path, object_file_path, *shared_libs] + print(cmd) + subprocess.run(cmd, check=True) + print(f"Successfully created shared library: {lib_path}") + + +if __name__ == "__main__": + main() 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 new file mode 100644 index 00000000..27883585 --- /dev/null +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp @@ -0,0 +1,97 @@ +// clang-format off +/* + * SPDX-FileCopyrightText: Copyright (c) 2023 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. + */ + +// clang-format on +// This example shows how to interface with an AOT compiled function in a C++ +// bundle. to build and run the example, run the following command in project +// root bash +// examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh + +#include +#include +#include +#include +#include +#include + +namespace ffi = tvm::ffi; + +struct CUDANDAlloc { + void AllocData(DLTensor *tensor) { + size_t data_size = ffi::GetDataSize(*tensor); + void *ptr = nullptr; + cudaError_t err = cudaMalloc(&ptr, data_size); + TVM_FFI_ICHECK_EQ(err, cudaSuccess) + << "cudaMalloc failed: " << cudaGetErrorString(err); + tensor->data = ptr; + } + + void FreeData(DLTensor *tensor) { + if (tensor->data != nullptr) { + cudaError_t err = cudaFree(tensor->data); + TVM_FFI_ICHECK_EQ(err, cudaSuccess) + << "cudaFree failed: " << cudaGetErrorString(err); + tensor->data = nullptr; + } + } +}; + +inline ffi::Tensor Empty(ffi::Shape shape, DLDataType dtype, DLDevice device) { + return ffi::Tensor::FromNDAlloc(CUDANDAlloc(), shape, dtype, device); +} + +// symbol from the shared library +extern "C" int __tvm_ffi_add_one(void *, const TVMFFIAny *, int32_t, + TVMFFIAny *); + +// Redirects into the exported function in object +void CallAddOne(ffi::TensorView x, ffi::TensorView y) { + tvm::ffi::Function::InvokeExternC(nullptr, __tvm_ffi_add_one, x, y); +} + +int main() { + DLDataType f32_dtype{kDLFloat, 32, 1}; + DLDevice cuda_device{kDLCUDA, 0}; + + constexpr int ARRAY_SIZE = 10; + + ffi::Tensor x = Empty({ARRAY_SIZE}, f32_dtype, cuda_device); + ffi::Tensor y = Empty({ARRAY_SIZE}, f32_dtype, cuda_device); + + std::vector host_x(ARRAY_SIZE); + for (int i = 0; i < ARRAY_SIZE; ++i) { + host_x[i] = static_cast(i); + } + + size_t nbytes = host_x.size() * sizeof(float); + cudaError_t err = + cudaMemcpy(x.data_ptr(), host_x.data(), nbytes, cudaMemcpyHostToDevice); + TVM_FFI_ICHECK_EQ(err, cudaSuccess) + << "cudaMemcpy host to device failed: " << cudaGetErrorString(err); + + // Call into the FFI function; tensors remain on device because they carry a + // kDLCUDA device tag. + CallAddOne(x, y); + + std::vector host_y(host_x.size()); + err = cudaMemcpy(host_y.data(), y.data_ptr(), nbytes, cudaMemcpyDeviceToHost); + TVM_FFI_ICHECK_EQ(err, cudaSuccess) + << "cudaMemcpy device to host failed: " << cudaGetErrorString(err); + + std::cout << "y after add_one_cuda(x, y)" << std::endl; + for (float value : host_y) { + std::cout << value << " "; + } + std::cout << std::endl; + return 0; +} 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 new file mode 100755 index 00000000..e81945ec --- /dev/null +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh @@ -0,0 +1,48 @@ +# 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. + +#!/bin/bash +CUDA_DIALECT_PATH="build/lib/" +export LD_LIBRARY_PATH=${CUDA_DIALECT_PATH}:`tvm-ffi-config --libdir` + +CUDA_HOME=/usr/local/cuda +SOURCE_FILE="$(dirname "$0")/aot_use_in_cpp_bundle.cpp" + +echo "Compiling the executable..." +g++ -o build/aot_use_in_cpp_bundle \ + -I${CUDA_HOME}/include \ + `tvm-ffi-config --cxxflags` \ + ${SOURCE_FILE} build/add_one.o \ + -L${CUDA_DIALECT_PATH} \ + -L${CUDA_HOME}/lib64 \ + -lcuda_dialect_runtime -lcuda -lcudart \ + `tvm-ffi-config --ldflags` \ + `tvm-ffi-config --libs` + +echo "Running the executable..." +./build/aot_use_in_cpp_bundle diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_jax.py b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_jax.py new file mode 100644 index 00000000..dd1f566a --- /dev/null +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_jax.py @@ -0,0 +1,52 @@ +# 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 jax +import jax.numpy as jnp +import jax_tvm_ffi +import cutlass.cute as cute + + +# now load it back +def main(): + a_jax = jnp.arange(10, dtype=jnp.float32) + b_jax = jnp.zeros(10, dtype=jnp.float32) + lib_path = "./build/add_one.so" + aot_mod = cute.runtime.load_module(lib_path) + jax_tvm_ffi.register_ffi_target("add_one_cute", aot_mod.add_one, platform="gpu") + b_jax = jax.ffi.ffi_call( + "add_one_cute", + jax.ShapeDtypeStruct(a_jax.shape, a_jax.dtype), + vmap_method="broadcast_all", + )(a_jax) + print("result of b after aot_mod.add_one(a, b)") + print(b_jax) + + +if __name__ == "__main__": + main() diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_torch.py b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_torch.py new file mode 100644 index 00000000..0f0fc046 --- /dev/null +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_torch.py @@ -0,0 +1,44 @@ +# 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.cute as cute +import torch + +# now load it back +def main(): + a_torch = torch.arange(10, dtype=torch.float32, device="cuda") + b_torch = torch.zeros(10, dtype=torch.float32, device="cuda") + lib_path = "./build/add_one.so" + aot_mod = cute.runtime.load_module(lib_path) + aot_mod.add_one(a_torch, b_torch) + print("result of b after aot_mod.add_one(a, b)") + print(b_torch) + + +if __name__ == "__main__": + main() diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/error_reporting.py b/examples/python/CuTeDSL/cute/tvm_ffi/error_reporting.py new file mode 100644 index 00000000..018a9eb2 --- /dev/null +++ b/examples/python/CuTeDSL/cute/tvm_ffi/error_reporting.py @@ -0,0 +1,71 @@ +# 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. + +"""Example demonstrating how to use TVM-FFI ABI with CuTe. + +This example shows how to: +1. Compile a CuTe function with "--enable-tvm-ffi" option +2. Directly use the compiled function to work with torch.Tensor + +To run this example: + +.. code-block:: bash + + python examples/cute/tvm_ffi/error_reporting.py +""" +import torch +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack + + +@cute.kernel +def device_add_one(a: cute.Tensor, b: cute.Tensor): + for i in range(a.shape[0]): + b[i] = a[i] + 1 + + +@cute.jit +def add_one(a: cute.Tensor, b: cute.Tensor): + """b = a + 1""" + device_add_one(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1)) + + +def main(): + # compile the kernel with "--enable-tvm-ffi" option + a_torch = torch.arange(10, dtype=torch.float32, device="cuda") + b_torch = torch.zeros(10, dtype=torch.float32, device="cuda") + a_cute = from_dlpack(a_torch, enable_tvm_ffi=True) + b_cute = from_dlpack(b_torch, enable_tvm_ffi=True) + # compile the kernel with "--enable-tvm-ffi" option + compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi") + # should raise an error because of shape mismatch + compiled_add_one(torch.arange(5, dtype=torch.float32, device="cuda"), b_cute) + + +if __name__ == "__main__": + main() diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_jax.py b/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_jax.py new file mode 100644 index 00000000..62b2abf8 --- /dev/null +++ b/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_jax.py @@ -0,0 +1,93 @@ +# 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. + +"""Example demonstrating how to use TVM-FFI ABI with CuTe. + +This example shows how to: +1. Compile a CuTe function with "--enable-tvm-ffi" option +2. Directly use the compiled function to work with JAX + +To run this example: + +.. code-block:: bash + + pip install jax-tvm-ffi + pip install jax[cuda13] + + python examples/cute/tvm_ffi/jit_and_use_in_jax.py +""" + +import jax +from jax import numpy as jnp +import jax_tvm_ffi +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack + + +@cute.kernel +def device_add_one(a: cute.Tensor, b: cute.Tensor): + for i in range(a.shape[0]): + b[i] = a[i] + 1 + + +@cute.jit +def add_one(a: cute.Tensor, b: cute.Tensor): + """b = a + 1""" + device_add_one(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1)) + + +def main(): + # compile the kernel with "--enable-tvm-ffi" option + a_jax = jnp.arange( + 10, + dtype=jnp.float32, + ) + b_jax = jnp.zeros( + 10, + dtype=jnp.float32, + ) + a_cute = from_dlpack(a_jax, enable_tvm_ffi=True).mark_layout_dynamic() + b_cute = from_dlpack(b_jax, enable_tvm_ffi=True).mark_layout_dynamic() + # compile the kernel with "--enable-tvm-ffi" option + compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi") + + # register the compiled function to JAX as a FFI target + jax_tvm_ffi.register_ffi_target("add_one_cute", compiled_add_one, platform="gpu") + a_jax = jnp.arange(10, dtype=jnp.float32) + # call the compiled function using JAX FFI + b_jax = jax.ffi.ffi_call( + "add_one_cute", + jax.ShapeDtypeStruct(a_jax.shape, a_jax.dtype), + vmap_method="broadcast_all", + )(a_jax) + print("result of b_jax after add_one_cute") + print(b_jax) + + +if __name__ == "__main__": + main() diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_torch.py b/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_torch.py new file mode 100644 index 00000000..bc257701 --- /dev/null +++ b/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_torch.py @@ -0,0 +1,83 @@ +# 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. + +"""Example demonstrating how to use TVM-FFI ABI with CuTe. + +This example shows how to: +1. Compile a CuTe function with "--enable-tvm-ffi" option +2. Directly use the compiled function to work with torch.Tensor + +To run this example: + +.. code-block:: bash + + python examples/cute/tvm_ffi/jit_and_use_in_torch.py +""" +import torch +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack + + +@cute.kernel +def device_add_one(a: cute.Tensor, b: cute.Tensor): + for i in range(a.shape[0]): + b[i] = a[i] + 1 + + +@cute.jit +def add_one(a: cute.Tensor, b: cute.Tensor): + """b = a + 1""" + device_add_one(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1)) + + +def main(): + # compile the kernel with "--enable-tvm-ffi" option + a_torch = torch.arange(10, dtype=torch.float32, device="cuda") + b_torch = torch.zeros(10, dtype=torch.float32, device="cuda") + a_cute = from_dlpack(a_torch, enable_tvm_ffi=True).mark_layout_dynamic() + b_cute = from_dlpack(b_torch, enable_tvm_ffi=True).mark_layout_dynamic() + # compile the kernel with "--enable-tvm-ffi" option + compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi") + # run the compiled function by passing in cute.Tensor as input + # you need to set enable_tvm_ffi=True for now + compiled_add_one(a_cute, b_cute) + # print the result + print("result of b after compiled_add_one(a, b)") + print(b_torch) + a_torch = a_torch + 1 + # We can directly pass in torch.Tensor as input + # the call overhead is optimized so it is very fast to pass in torch.Tensor as input + # takes about less than 0.5us per call likely in terms of API overhead + compiled_add_one(a_torch, b_torch) + # print the result + print("result of b after compiled_add_one(a, b)") + print(b_torch) + + +if __name__ == "__main__": + main() diff --git a/examples/python/CuTeDSL/helpers/__init__.py b/examples/python/CuTeDSL/helpers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/python/CuTeDSL/helpers/fmha_helpers.py b/examples/python/CuTeDSL/helpers/fmha_helpers.py new file mode 100644 index 00000000..6d49a9f1 --- /dev/null +++ b/examples/python/CuTeDSL/helpers/fmha_helpers.py @@ -0,0 +1,975 @@ +# 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 +# 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. + +import enum +from typing import Tuple, Optional +import cutlass +from cutlass.cute.typing import Boolean + +from cutlass.cutlass_dsl import ( + Int32, + Float32, + min, + extract_mlir_values, + new_from_mlir_values, +) +from cutlass.utils.hardware_info import HardwareInfo +from cutlass.utils import WorkTileInfo +import cutlass.cute as cute + +############################################################################## +# Fmha static tile scheduler +############################################################################## + + +class FmhaStaticTileSchedulerParams: + """A class to represent parameters for the FMHA (Fused Multi-Head Attention) static tile scheduler. + + This class holds the configuration parameters needed to initialize and configure + the tile scheduler for FMHA operations. + + :ivar is_persistent: Whether to use persistent kernel mode. + :type is_persistent: bool + :ivar problem_shape_mbh: Problem shape in (M, B, H) format. + :type problem_shape_mbh: cute.Shape + """ + + def __init__( + self, + is_persistent: bool, + problem_shape_mbh: cute.Shape, + *, + loc=None, + ip=None, + ): + """ + Initializes the FmhaStaticTileSchedulerParams with the given parameters. + + :param is_persistent: Whether to use persistent kernel mode. + :type is_persistent: bool + :param problem_shape_mbh: Problem shape in (M, B, H) format. + :type problem_shape_mbh: cute.Shape + """ + self.is_persistent = is_persistent + self.problem_shape_mbh = problem_shape_mbh + self._loc = loc + self._ip = ip + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [self.problem_shape_mbh]: + obj_values = extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + for obj, n_items in zip([self.problem_shape_mbh], self._values_pos): + obj_list.append(new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return FmhaStaticTileSchedulerParams( + self.is_persistent, *(tuple(obj_list)), loc=self._loc + ) + + +class FmhaStaticTileScheduler: + """A static tile scheduler for FMHA (Fused Multi-Head Attention) operations. + + This class manages the scheduling of work tiles for FMHA kernels, supporting + both persistent and non-persistent kernel modes. It tracks the current work + position and advances through the problem space efficiently. + + :ivar _params: Scheduler parameters. + :type _params: FmhaStaticTileSchedulerParams + :ivar _blk_coord: Block coordinates. + :type _blk_coord: cute.Coord + :ivar _grid_shape: Grid shape for the kernel. + :type _grid_shape: cute.Shape + :ivar _is_persistent: Whether to use persistent kernel mode. + :type _is_persistent: bool + :ivar _current_work_linear_idx: Current linear work index. + :type _current_work_linear_idx: Int32 + :ivar _problem_shape_mbh: Problem shape in (M, B, H) format. + :type _problem_shape_mbh: cute.Layout + :ivar _num_blocks: Number of blocks in the problem. + :type _num_blocks: Int32 + :ivar _is_first_block: Whether this is the first block. + :type _is_first_block: bool + :ivar num_persistent_sm: Number of persistent SMs. + :type num_persistent_sm: Int32 + """ + + def __init__( + self, + params: FmhaStaticTileSchedulerParams, + current_work_linear_idx: Int32, + blk_coord: cute.Coord, + grid_shape: cute.Shape, + *, + loc=None, + ip=None, + ): + """ + Initializes the FmhaStaticTileScheduler with the given parameters. + + :param params: Scheduler parameters. + :type params: FmhaStaticTileSchedulerParams + :param current_work_linear_idx: Current linear work index. + :type current_work_linear_idx: Int32 + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param grid_shape: Grid shape for the kernel. + :type grid_shape: cute.Shape + """ + self._params = params + self._blk_coord = blk_coord + self._grid_shape = grid_shape + self._is_persistent = params.is_persistent + self._current_work_linear_idx = current_work_linear_idx + self._problem_shape_mbh = cute.make_layout( + params.problem_shape_mbh, loc=loc, ip=ip + ) + self._num_blocks = cute.size(self._problem_shape_mbh, loc=loc, ip=ip) + self._is_first_block = True + self.num_persistent_sm = cute.size(grid_shape, loc=loc, ip=ip) + self._loc = loc + self._ip = ip + + # called by host + @staticmethod + def get_grid_shape( + params: FmhaStaticTileSchedulerParams, + *, + loc=None, + ip=None, + ) -> cute.Shape: + """ + Determine the grid shape for the FMHA kernel. + + For persistent kernels, the grid shape is limited by the number of SMs + (Streaming Multiprocessors) available on the device. For non-persistent + kernels, the grid shape matches the problem shape. + + :param params: Scheduler parameters. + :type params: FmhaStaticTileSchedulerParams + + :return: Grid shape as (M, B, H) tuple. + :rtype: cute.Shape + """ + if params.is_persistent: + hardware_info = HardwareInfo() + sm_count = hardware_info.get_device_multiprocessor_count() + return ( + min(sm_count, cute.size(params.problem_shape_mbh, loc=loc, ip=ip)), + 1, + 1, + ) + else: + return params.problem_shape_mbh + + @staticmethod + def check_valid_work_for_seqlen_q( + q_tiler: int, + current_idx: Int32, + seqlen_q: Int32, + ) -> Boolean: + """ + Check if the current work index is valid for the given query sequence length. + + This method verifies that the current work tile index multiplied by the + query tiler size is within the bounds of the query sequence length. + + :param q_tiler: Query tiler size. + :type q_tiler: int + :param current_idx: Current work index. + :type current_idx: Int32 + :param seqlen_q: Query sequence length. + :type seqlen_q: Int32 + + :return: True if the work is valid, False otherwise. + :rtype: Boolean + """ + return current_idx * q_tiler < seqlen_q + + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + """ + Get information about the current work tile. + + Determines if the current work is valid and computes the tile coordinates + based on whether the kernel is persistent or non-persistent. + + :return: WorkTileInfo containing tile coordinates and validity flag. + :rtype: WorkTileInfo + """ + is_valid = ( + self._current_work_linear_idx < self._num_blocks + if self._is_persistent + else self._is_first_block + ) + + blk_coord = (0, 0, 0) + if self._is_persistent: + blk_coord = self._problem_shape_mbh.get_hier_coord( + self._current_work_linear_idx, loc=loc, ip=ip + ) + else: + blk_coord = self._blk_coord + + # cur_tile_coord is (mid, 0, (bid, hid)) + cur_tile_coord = ( + blk_coord[0], + 0, + (blk_coord[1], blk_coord[2]), + ) + + return WorkTileInfo(cur_tile_coord, is_valid) + + def initial_work_tile_info(self, *, loc=None, ip=None): + """ + Get the initial work tile information. + + :return: Initial WorkTileInfo. + :rtype: WorkTileInfo + """ + return self.get_current_work(loc=loc, ip=ip) + + def advance_to_next_work(self, *, advance_count=1, loc=None, ip=None): + """ + Advance to the next work tile. + + For persistent kernels, advances by the number of persistent SMs. + For non-persistent kernels, marks that the first block has been processed. + + :param advance_count: Number of steps to advance (default: 1). + :type advance_count: int + """ + if self._is_persistent: + self._current_work_linear_idx += advance_count * self.num_persistent_sm + self._is_first_block = False + + def __extract_mlir_values__(self): + values = extract_mlir_values(self._params) + values.extend(extract_mlir_values(self._current_work_linear_idx)) + values.extend(extract_mlir_values(self._blk_coord)) + values.extend(extract_mlir_values(self._grid_shape)) + return values + + def __new_from_mlir_values__(self, values): + assert len(values) == 10 + new_params = new_from_mlir_values(self._params, values[0:3]) + new_current_work_linear_idx = new_from_mlir_values( + self._current_work_linear_idx, [values[3]] + ) + new_blk_coord = new_from_mlir_values(self._blk_coord, values[4:7]) + new_grid_shape = new_from_mlir_values(self._grid_shape, values[7:]) + return FmhaStaticTileScheduler( + new_params, new_current_work_linear_idx, new_blk_coord, new_grid_shape + ) + + +def create_fmha_static_tile_scheduler( + params: FmhaStaticTileSchedulerParams, + blk_coord: cute.Coord, + grid_shape: cute.Shape, +) -> FmhaStaticTileScheduler: + """ + Create a new FMHA static tile scheduler. + + :param params: Scheduler parameters. + :type params: FmhaStaticTileSchedulerParams + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param grid_shape: Grid shape. + :type grid_shape: cute.Shape + + :return: New FmhaStaticTileScheduler instance. + :rtype: FmhaStaticTileScheduler + """ + return FmhaStaticTileScheduler(params, blk_coord[0], blk_coord, grid_shape) + + +def create_fmha_static_tile_scheduler_params( + is_persistent: bool, + problem_shape_mbh: cute.Shape, +) -> FmhaStaticTileSchedulerParams: + """ + Create FMHA static tile scheduler parameters. + + :param is_persistent: Whether to use persistent kernel mode. + :type is_persistent: bool + :param problem_shape_mbh: Problem shape in (M, B, H) format. + :type problem_shape_mbh: cute.Shape + + :return: New FmhaStaticTileSchedulerParams instance. + :rtype: FmhaStaticTileSchedulerParams + """ + return FmhaStaticTileSchedulerParams(is_persistent, problem_shape_mbh) + + +def compute_grid( + o_shape: cute.Shape, + cta_tiler: Tuple[int, int, int], + is_persistent: bool, +) -> Tuple[FmhaStaticTileSchedulerParams, Tuple[int, int, int]]: + """ + Compute grid parameters for FMHA operation. + + This function calculates the appropriate grid shape and scheduler parameters + based on the output tensor shape, CTA (Cooperative Thread Array) tiler, + and whether to use persistent kernel mode. + + The output tensor o has shape (s, d, ((h_r, h_k), b)) where: + - s: sequence length + - d: head dimension + - h_r: number of heads for query + - h_k: number of heads for key + - b: batch size + + :param o_shape: Output tensor shape for grid computation. + :type o_shape: cute.Shape + :param cta_tiler: CTA tiler dimensions (M, N, K). + :type cta_tiler: Tuple[int, int, int] + :param is_persistent: Whether to use persistent kernel mode. + :type is_persistent: bool + + :return: Tuple of (scheduler_params, grid_shape). + :rtype: Tuple[FmhaStaticTileSchedulerParams, Tuple[int, int, int]] + """ + tile_sched_params = create_fmha_static_tile_scheduler_params( + is_persistent, + ( + cute.ceil_div(cute.size(o_shape[0]), cta_tiler[0]), + cute.size(o_shape[2][0]), + cute.size(o_shape[2][1]), + ), + ) + grid = FmhaStaticTileScheduler.get_grid_shape(tile_sched_params) + + return tile_sched_params, grid + + +############################################################################## +# Fused Mask +############################################################################## + + +class MaskEnum(enum.Enum): + """Enumeration of mask types for FMHA operations. + + - RESIDUAL_MASK: Residual mask for handling variable sequence lengths + - WINDOW_MASK: Window mask for attention which also includes causal and no mask + - WINDOW_MASK_INFERENCE: Same as the window mask, but has the limitation that the end of q is aligned with the end of k + - WINDOW_MASK_BWD: Window mask for backward pass + - WINDOW_MASK_BWD_INFERENCE: Same as the window mask for backward pass, but has the limitation that the end of q is aligned with the end of k + """ + + RESIDUAL_MASK = enum.auto() + RESIDUAL_MASK_BWD = enum.auto() + WINDOW_MASK = enum.auto() + WINDOW_MASK_INFERENCE = enum.auto() + WINDOW_MASK_BWD = enum.auto() + WINDOW_MASK_BWD_INFERENCE = enum.auto() + + +class FusedMask: + """A fused mask implementation for FMHA operations. + + This class handles different types of attention masks including no mask, + residual mask for variable sequence lengths, and causal mask for + autoregressive attention patterns. + + The class provides methods to: + - Calculate trip counts for different mask types + - Apply masks to attention scores + - Handle masked and unmasked trip calculations + """ + + def get_trip_count( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Calculate the number of trips needed for the current block. + + The trip count depends on the mask type and the block coordinates. + For causal masks, it considers the autoregressive constraint. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Number of trips needed. + :rtype: Int32 + """ + result = 0 + offset = 0 + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + if cutlass.const_expr(mask_type == MaskEnum.RESIDUAL_MASK): + result = cute.ceil_div(seqlen_k, tile_shape[1]) + if cutlass.const_expr(mask_type is MaskEnum.RESIDUAL_MASK_BWD): + result = cute.ceil_div(seqlen_q, tile_shape[0]) + if cutlass.const_expr( + mask_type == MaskEnum.WINDOW_MASK + or mask_type == MaskEnum.WINDOW_MASK_INFERENCE + ): + if cutlass.const_expr(window_size_right is None): + result = cute.ceil_div(seqlen_k, tile_shape[1]) + else: + max_idx_q = (blk_coord[0] + 1) * tile_shape[0] + idx_k = max_idx_q + offset + window_size_right + tmp_blocks_k = cute.ceil_div(idx_k, tile_shape[1]) + max_blocks_k = cute.ceil_div(seqlen_k, tile_shape[1]) + result = min(max_blocks_k, tmp_blocks_k) + if cutlass.const_expr( + mask_type == MaskEnum.WINDOW_MASK_BWD + or mask_type == MaskEnum.WINDOW_MASK_BWD_INFERENCE + ): + if cutlass.const_expr(window_size_left is None): + result = cute.ceil_div(seqlen_q, tile_shape[0]) + else: + max_idx_k = (blk_coord[1] + 1) * tile_shape[1] + idx_k = max_idx_k + offset + window_size_left + tmp_blocks_q = cute.ceil_div(idx_k, tile_shape[0]) + max_blocks_q = cute.ceil_div(seqlen_q, tile_shape[0]) + result = min(max_blocks_q, tmp_blocks_q) + start_block = FusedMask.get_trip_start( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + result = result - start_block + return result + + @cute.jit + def get_trip_start( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Get the start of the trip for the current block. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + """ + result = 0 + offset = 0 + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + if cutlass.const_expr( + mask_type is MaskEnum.WINDOW_MASK + or mask_type is MaskEnum.WINDOW_MASK_INFERENCE + ): + if cutlass.const_expr(window_size_left is not None): + min_idx_q = blk_coord[0] * tile_shape[0] + idx_k = min_idx_q + offset - window_size_left + tmp_blocks_k = idx_k // tile_shape[1] + result = max(tmp_blocks_k, result) + if cutlass.const_expr( + mask_type is MaskEnum.WINDOW_MASK_BWD + or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE + ): + if cutlass.const_expr(window_size_right is not None): + min_idx_k = blk_coord[1] * tile_shape[1] + idx_q = min_idx_k + offset - window_size_right + tmp_blocks_q = idx_q // tile_shape[0] + result = max(tmp_blocks_q, result) + return result + + @cute.jit + def get_leading_mask_id( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Tuple[Int32, Int32]: + """ + Get the begin and end tile idx for the leading mask. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Tuple of (begin, end) tile idx for the leading mask. + :rtype: Tuple[Int32, Int32] + """ + offset = 0 + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + leading_mask_begin = FusedMask.get_trip_start( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + trip_count = FusedMask.get_trip_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + + leading_mask_end = leading_mask_begin + if cutlass.const_expr( + mask_type is MaskEnum.WINDOW_MASK + or mask_type is MaskEnum.WINDOW_MASK_INFERENCE + ): + if cutlass.const_expr(window_size_left is not None): + min_idx_q = ( + (blk_coord[0] + 1) * tile_shape[0] + offset - window_size_left + ) + leading_mask_end = min( + cute.ceil_div(min_idx_q, tile_shape[1]) - 1, + trip_count + leading_mask_begin - 1, + ) + else: + leading_mask_end = leading_mask_begin - 1 + elif cutlass.const_expr( + mask_type is MaskEnum.WINDOW_MASK_BWD + or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE + ): + if cutlass.const_expr(window_size_right is not None): + min_idx_k = ( + (blk_coord[1] + 1) * tile_shape[1] + offset - window_size_right + ) + leading_mask_end = cute.ceil_div(min_idx_k, tile_shape[0]) - 1 + else: + leading_mask_end = leading_mask_begin - 1 + return leading_mask_begin, leading_mask_end + + @cute.jit + def get_trailing_mask_id( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Tuple[Optional[Int32], Optional[Int32]]: + """ + Get the begin and end tile idx for the trailing mask. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Tuple of (begin, end) tile idx for the trailing mask. + :rtype: Tuple[Int32, Int32] + """ + offset = 0 + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + trip_start = FusedMask.get_trip_start( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + trip_count = FusedMask.get_trip_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + + trailing_mask_begin, trailing_mask_end = None, None + if cutlass.const_expr( + mask_type is MaskEnum.WINDOW_MASK + or mask_type is MaskEnum.WINDOW_MASK_INFERENCE + ): + if cutlass.const_expr(window_size_right is not None): + min_idx_q = blk_coord[0] * tile_shape[0] + offset + window_size_right + trailing_mask_begin = min( + min_idx_q // tile_shape[1], trip_count + trip_start - 1 + ) + trailing_mask_end = trip_count + trip_start - 1 + else: + # last tile, we always apply mask on it regardless whether it's a residual tile + trailing_mask_begin = trip_count + trip_start - 1 + trailing_mask_end = trip_count + trip_start - 1 + else: + if cutlass.const_expr(window_size_left is not None): + min_idx_k = blk_coord[1] * tile_shape[1] + offset + window_size_left + 1 + max_idx_k = ( + (blk_coord[1] + 1) * tile_shape[1] + offset + window_size_left + ) + trailing_mask_begin = min( + cute.ceil_div(min_idx_k, tile_shape[0]) - 1, + trip_count + trip_start - 1, + ) + trailing_mask_end = min( + cute.ceil_div(max_idx_k, tile_shape[0]) - 1, + trip_count + trip_start - 1, + ) + else: + # last tile, we always apply mask on it regardless whether it's a residual tile + trailing_mask_begin = trip_count + trip_start - 1 + trailing_mask_end = trip_count + trip_start - 1 + + return trailing_mask_begin, trailing_mask_end + + @cute.jit + def get_masked_leading_count( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Calculate the number of masked trips for the leading mask. + + This is used for blocks that need special handling due to masking. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Number of masked trips. + :rtype: Int32 + """ + result = 0 + if cutlass.const_expr( + mask_type is not MaskEnum.RESIDUAL_MASK + and mask_type is not MaskEnum.RESIDUAL_MASK_BWD + ): + if cutlass.const_expr( + window_size_left is not None or window_size_right is not None + ): + leading_mask_begin, leading_mask_end = FusedMask.get_leading_mask_id( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + result = max(leading_mask_end - leading_mask_begin + 1, 0) + + return result + + @cute.jit + def get_masked_trailing_count( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + rem_count: Optional[Int32] = 0, + ) -> Int32: + """ + Calculate the number of masked trips for the trailing mask. + + This is used for blocks that need special handling due to masking. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + :param rem_count: Remaining count from previous calculations. + :type rem_count: Int32 + + :return: Number of masked trips. + :rtype: Int32 + """ + result = 0 + + if cutlass.const_expr( + mask_type is not MaskEnum.RESIDUAL_MASK + and mask_type is not MaskEnum.RESIDUAL_MASK_BWD + ): + if cutlass.const_expr( + window_size_left is not None or window_size_right is not None + ): + trailing_mask_begin, trailing_mask_end = FusedMask.get_trailing_mask_id( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + leading_mask_begin, leading_mask_end = FusedMask.get_leading_mask_id( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + if cutlass.const_expr( + trailing_mask_begin is not None and trailing_mask_end is not None + ): + if trailing_mask_begin <= leading_mask_end: + result = max(trailing_mask_end - leading_mask_end, 0) + else: + result = max(trailing_mask_end - trailing_mask_begin + 1, 0) + else: + if seqlen_k % tile_shape[1] != 0: + result = 1 + else: + result = 0 + + return result + rem_count + + @cute.jit + def get_unmasked_trip_count( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Calculate the number of unmasked trips for the current block. + + This represents the number of trips that don't require special + masking treatment. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Number of unmasked trips. + :rtype: Int32 + """ + result = ( + FusedMask.get_trip_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + - FusedMask.get_masked_leading_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + - FusedMask.get_masked_trailing_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + 0, + ) + ) + return result + + @cute.jit + def apply_mask( + mask_type: MaskEnum, + acc_qk: cute.Tensor, + index_qk: cute.Tensor, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[int] = None, + window_size_right: Optional[int] = None, + index_transform: cutlass.Constexpr = lambda index_q, index_k: ( + index_q, + index_k, + ), + ): + """ + Apply the appropriate mask to the attention scores. + + This method modifies the attention scores (acc_qk) based on the mask type + and the positions in the index tensor. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param acc_qk: Accumulated QK attention scores tensor. + :type acc_qk: cute.Tensor + :param index_qk: Index tensor containing position information. + :type index_qk: cute.Tensor + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Optional[int] + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[int] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[int] + """ + + tidx, tidy, tidx = cute.arch.thread_idx() + offset = 0 + offset = ( + seqlen_k - seqlen_q + if cutlass.const_expr( + mask_type is MaskEnum.WINDOW_MASK_INFERENCE + or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE + ) + else 0 + ) + for i in cutlass.range_constexpr(cute.size(acc_qk)): + index_q, index_k = index_transform(*index_qk[i]) + if cutlass.const_expr( + window_size_left is not None or window_size_right is not None + ): + if cutlass.const_expr(window_size_left is None): + if index_q + offset + window_size_right < index_k: + acc_qk[i] = -Float32.inf + if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask + acc_qk[i] = -Float32.inf + elif cutlass.const_expr(window_size_right is None): + if index_q + offset - window_size_left > index_k: + acc_qk[i] = -Float32.inf + if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask + acc_qk[i] = -Float32.inf + else: + max_K_index = min(index_q + offset + window_size_right, seqlen_k) + min_K_index = max(0, index_q + offset - window_size_left) + if index_k > max_K_index or index_k < min_K_index: + acc_qk[i] = -Float32.inf + if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask + acc_qk[i] = -Float32.inf + + if cutlass.const_expr( + mask_type == MaskEnum.RESIDUAL_MASK + or mask_type == MaskEnum.RESIDUAL_MASK_BWD + ): + if index_k >= seqlen_k or index_q >= seqlen_q: + acc_qk[i] = -Float32.inf diff --git a/examples/python/CuTeDSL/helpers/sparse_utils.py b/examples/python/CuTeDSL/helpers/sparse_utils.py new file mode 100644 index 00000000..24b3f791 --- /dev/null +++ b/examples/python/CuTeDSL/helpers/sparse_utils.py @@ -0,0 +1,457 @@ +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 new file mode 100644 index 00000000..3264f191 --- /dev/null +++ b/examples/python/CuTeDSL/helpers/test_sparse_utils.py @@ -0,0 +1,104 @@ +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 ce7afa45..ed960de7 100644 --- a/examples/python/CuTeDSL/hopper/dense_gemm.py +++ b/examples/python/CuTeDSL/hopper/dense_gemm.py @@ -38,6 +38,7 @@ 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 @@ -624,11 +625,11 @@ class HopperWgmmaGemmKernel: consumer_group=mainloop_pipeline_consumer_group, tx_count=tma_copy_bytes, cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, ) # Cluster arrive after barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # /////////////////////////////////////////////////////////////////////////////// # Generate smem tensor A/B @@ -717,10 +718,7 @@ class HopperWgmmaGemmKernel: # Cluster wait # /////////////////////////////////////////////////////////////////////////////// # cluster wait for barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_wait() - else: - cute.arch.sync_threads() + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) # ///////////////////////////////////////////////////////////////////////////// # Prefetch # ///////////////////////////////////////////////////////////////////////////// diff --git a/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py b/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py index cd6e2d6a..9f9aec2b 100644 --- a/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py +++ b/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py @@ -37,6 +37,7 @@ 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 @@ -634,11 +635,11 @@ class HopperWgmmaGemmPersistentKernel: consumer_group=mainloop_pipeline_consumer_group, tx_count=tma_copy_bytes, cta_layout_vmnk=cute.make_layout((1, *cta_layout_mnk.shape)), + defer_sync=True, ) # Cluster arrive after barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # Generate smem tensor A/B sA = storage.sA.get_tensor( @@ -718,10 +719,7 @@ class HopperWgmmaGemmPersistentKernel: k_tile_cnt = cute.size(gA_mkl, mode=[3]) # Cluster wait for barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_wait() - else: - cute.arch.sync_threads() + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) is_dma_warp_group = warp_group_idx < self.num_dma_warp_groups if is_dma_warp_group: diff --git a/examples/python/CuTeDSL/hopper/fmha.py b/examples/python/CuTeDSL/hopper/fmha.py index 0d04831d..c9a94091 100644 --- a/examples/python/CuTeDSL/hopper/fmha.py +++ b/examples/python/CuTeDSL/hopper/fmha.py @@ -95,15 +95,18 @@ import cutlass.cute.testing as testing 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 from cutlass.cute.runtime import from_dlpack -current_dir = os.path.dirname(os.path.abspath(__file__)) -sys.path.insert(0, os.path.join(current_dir, "..")) -from utils import fmha_helpers as fmha_utils +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 class HopperFusedMultiHeadAttentionForward: @@ -648,11 +651,8 @@ class HopperFusedMultiHeadAttentionForward: # 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 - if cute.size(self.cluster_shape_mnk) > 1: - cute.arch.cluster_arrive_relaxed() - cute.arch.cluster_wait() - else: - cute.arch.sync_threads() + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True) + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk) if warp_idx == 0: cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q) @@ -1646,6 +1646,7 @@ class HopperFusedMultiHeadAttentionForward: producer_group=load_q_producer_group, consumer_group=load_q_consumer_group, tx_count=self.tma_copy_q_bytes, + defer_sync=True, ).make_participants() def make_and_init_load_kv_pipeline(self, load_kv_mbar_ptr): @@ -1663,6 +1664,7 @@ class HopperFusedMultiHeadAttentionForward: producer_group=load_kv_producer_group, consumer_group=load_kv_consumer_group, tx_count=self.tma_copy_kv_bytes, + defer_sync=True, ).make_participants() def make_and_init_tma_store_pipeline(self): @@ -1686,6 +1688,7 @@ class HopperFusedMultiHeadAttentionForward: pipeline.Agent.Thread, self.num_threads_per_warp_group, ), + defer_sync=True, ) @staticmethod diff --git a/include/cute/container/array.hpp b/include/cute/container/array.hpp index 31606e45..3fe2be77 100644 --- a/include/cute/container/array.hpp +++ b/include/cute/container/array.hpp @@ -447,13 +447,7 @@ struct tuple_element> namespace std { -#if defined(__CUDACC_RTC__) -template -struct tuple_size; - -template -struct tuple_element; -#endif +#include template struct tuple_size> diff --git a/include/cute/container/array_subbyte.hpp b/include/cute/container/array_subbyte.hpp index be6410b3..30cd438f 100644 --- a/include/cute/container/array_subbyte.hpp +++ b/include/cute/container/array_subbyte.hpp @@ -617,13 +617,7 @@ struct tuple_element> namespace std { -#if defined(__CUDACC_RTC__) -template -struct tuple_size; - -template -struct tuple_element; -#endif +#include template struct tuple_size> diff --git a/include/cute/container/tuple.hpp b/include/cute/container/tuple.hpp index e3dd6d27..71c7645c 100644 --- a/include/cute/container/tuple.hpp +++ b/include/cute/container/tuple.hpp @@ -701,13 +701,7 @@ struct tuple_element> namespace std { -#if defined(__CUDACC_RTC__) -template -struct tuple_size; - -template -struct tuple_element; -#endif +#include template struct tuple_size> diff --git a/include/cute/container/type_list.hpp b/include/cute/container/type_list.hpp index ff9498eb..2d2865f6 100644 --- a/include/cute/container/type_list.hpp +++ b/include/cute/container/type_list.hpp @@ -103,13 +103,7 @@ struct tuple_element> namespace std { -#if defined(__CUDACC_RTC__) -template -struct tuple_size; - -template -struct tuple_element; -#endif +#include template struct tuple_size> diff --git a/include/cute/numeric/arithmetic_tuple.hpp b/include/cute/numeric/arithmetic_tuple.hpp index 016ac5b6..8d4b7a91 100644 --- a/include/cute/numeric/arithmetic_tuple.hpp +++ b/include/cute/numeric/arithmetic_tuple.hpp @@ -513,13 +513,7 @@ struct tuple_element> namespace std { -#if defined(__CUDACC_RTC__) -template -struct tuple_size; - -template -struct tuple_element; -#endif +#include template struct tuple_size> diff --git a/include/cute/util/print.hpp b/include/cute/util/print.hpp index e6cc887a..5ce26fb0 100644 --- a/include/cute/util/print.hpp +++ b/include/cute/util/print.hpp @@ -127,6 +127,18 @@ print(uint4b_t a) { printf("%d", int(a)); } +CUTE_HOST_DEVICE +void +print(int6b_t a) { + printf("%d", int(a)); +} + +CUTE_HOST_DEVICE +void +print(uint6b_t a) { + printf("%d", int(a)); +} + CUTE_HOST_DEVICE void print(bin1_t a) { @@ -217,6 +229,16 @@ pretty_print(uint4b_t a) { printf("%*d", 5, int(a)); } +CUTE_HOST_DEVICE void +pretty_print(int6b_t a) { + printf("%*d", 5, int(a)); +} + +CUTE_HOST_DEVICE void +pretty_print(uint6b_t a) { + printf("%*d", 5, int(a)); +} + CUTE_HOST_DEVICE void pretty_print(bool v) { printf("%*d", 3, int(v)); diff --git a/include/cutlass/arch/barrier.h b/include/cutlass/arch/barrier.h index 2b0d4bb6..245e7d2d 100644 --- a/include/cutlass/arch/barrier.h +++ b/include/cutlass/arch/barrier.h @@ -283,7 +283,7 @@ class NamedBarrier { CUTLASS_DEVICE static void arrive_and_wait_internal(uint32_t num_threads, uint32_t barrier_id) { #if CUDA_BARRIER_ENABLED - asm volatile("bar.sync %0, %1;" : : "r"(barrier_id), "r"(num_threads)); + asm volatile("bar.sync %0, %1;" : : "r"(barrier_id), "r"(num_threads) : "memory"); cutlass::arch::synclog_emit_named_barrier_arrive_and_wait(__LINE__, num_threads, barrier_id); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -293,7 +293,7 @@ class NamedBarrier { CUTLASS_DEVICE static void arrive_and_wait_internal_unaligned(uint32_t num_threads, uint32_t barrier_id) { #if CUDA_BARRIER_ENABLED - asm volatile("barrier.sync %0, %1;" : : "r"(barrier_id), "r"(num_threads)); + asm volatile("barrier.sync %0, %1;" : : "r"(barrier_id), "r"(num_threads) : "memory"); cutlass::arch::synclog_emit_named_barrier_arrive_and_wait(__LINE__, num_threads, barrier_id); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -304,7 +304,7 @@ class NamedBarrier { static void arrive_internal(uint32_t num_threads, uint32_t barrier_id) { #if CUDA_BARRIER_ENABLED cutlass::arch::synclog_emit_named_barrier_arrive(__LINE__, num_threads, barrier_id); - asm volatile("bar.arrive %0, %1;" : : "r"(barrier_id), "r"(num_threads)); + asm volatile("bar.arrive %0, %1;" : : "r"(barrier_id), "r"(num_threads) : "memory"); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); #endif @@ -314,7 +314,7 @@ class NamedBarrier { static void arrive_internal_unaligned(uint32_t num_threads, uint32_t barrier_id) { #if CUDA_BARRIER_ENABLED cutlass::arch::synclog_emit_named_barrier_arrive(__LINE__, num_threads, barrier_id); - asm volatile("barrier.arrive %0, %1;" : : "r"(barrier_id), "r"(num_threads)); + asm volatile("barrier.arrive %0, %1;" : : "r"(barrier_id), "r"(num_threads) : "memory"); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); #endif @@ -396,7 +396,8 @@ public: "mbarrier.init.shared::cta.b64 [%1], %0; \n" "}" : - : "r"(arrive_count), "r"(smem_addr)); + : "r"(arrive_count), "r"(smem_addr) + : "memory"); cutlass::arch::synclog_emit_cluster_barrier_init(__LINE__, smem_addr, arrive_count); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -421,7 +422,8 @@ public: "DONE: \n\t" "}" : - : "r"(smem_addr), "r"(phase), "r"(ticks)); + : "r"(smem_addr), "r"(phase), "r"(ticks) + : "memory"); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -444,7 +446,8 @@ public: "selp.b32 %0, 1, 0, P1; \n\t" "}" : "=r"(waitComplete) - : "r"(smem_addr), "r"(phase), "r"(pred)); + : "r"(smem_addr), "r"(phase), "r"(pred) + : "memory"); return static_cast(waitComplete); #elif defined(__CUDA_ARCH__) @@ -467,7 +470,8 @@ public: "selp.b32 %0, 1, 0, P1; \n\t" "}" : "=r"(waitComplete) - : "r"(smem_addr), "r"(phase)); + : "r"(smem_addr), "r"(phase) + : "memory"); return static_cast(waitComplete); #elif defined(__CUDA_ARCH__) @@ -489,7 +493,8 @@ public: "mbarrier.arrive.shared::cluster.b64 _, [remAddr32];\n\t" "}" : - : "r"(smem_addr), "r"(cta_id)); + : "r"(smem_addr), "r"(cta_id) + : "memory"); } cutlass::arch::synclog_emit_cluster_barrier_arrive_cluster(__LINE__, smem_addr, cta_id, pred); @@ -508,7 +513,8 @@ public: "mbarrier.arrive.shared::cta.b64 _, [%0];\n\t" "}" : - : "r"(smem_addr)); + : "r"(smem_addr) + : "memory"); cutlass::arch::synclog_emit_cluster_barrier_arrive(__LINE__, smem_addr); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -524,7 +530,8 @@ public: "mbarrier.inval.shared::cta.b64 [%0]; \n\t" "}" : - : "r"(smem_addr)); + : "r"(smem_addr) + : "memory"); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); #endif @@ -585,7 +592,8 @@ struct ClusterTransactionBarrier : public ClusterBarrier { "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%1], %0; \n\t" "}" : - : "r"(transaction_bytes), "r"(smem_addr)); + : "r"(transaction_bytes), "r"(smem_addr) + : "memory"); cutlass::arch::synclog_emit_cluster_transaction_barrier_arrive_and_expect_tx(__LINE__, smem_addr, transaction_bytes); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -607,7 +615,8 @@ struct ClusterTransactionBarrier : public ClusterBarrier { "@p mbarrier.arrive.expect_tx.shared::cluster.b64 _, [remAddr32], %3;\n\t" "}" : - : "r"(smem_addr), "r"(cta_id), "r"(pred), "r"(transaction_bytes)); + : "r"(smem_addr), "r"(cta_id), "r"(pred), "r"(transaction_bytes) + : "memory"); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); #endif @@ -623,7 +632,8 @@ struct ClusterTransactionBarrier : public ClusterBarrier { "mbarrier.expect_tx.shared::cta.b64 [%1], %0; \n\t" "}" : - : "r"(transaction_bytes), "r"(smem_addr)); + : "r"(transaction_bytes), "r"(smem_addr) + : "memory"); cutlass::arch::synclog_emit_cluster_transaction_barrier_expect_transaction(__LINE__, smem_addr, transaction_bytes); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -644,7 +654,8 @@ struct ClusterTransactionBarrier : public ClusterBarrier { "@p mbarrier.complete_tx.shared::cluster.relaxed.cluster.b64 [%1], %0;" "}" : - : "r"(transaction_bytes), "r"(smem_addr), "r"(pred)); + : "r"(transaction_bytes), "r"(smem_addr), "r"(pred) + : "memory"); cutlass::arch::synclog_emit_cluster_transaction_barrier_complete_transaction(__LINE__, smem_addr, dst_cta_id, transaction_bytes, pred); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -704,7 +715,8 @@ void fence_barrier_init() { "{\n\t" "fence.mbarrier_init.release.cluster; \n" "}" - ::); + :: + : "memory"); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); #endif @@ -719,7 +731,8 @@ void fence_view_async_shared() { "{\n\t" "fence.proxy.async.shared::cta; \n" "}" - ::); + :: + : "memory"); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); #endif @@ -735,7 +748,8 @@ void cpasync_barrier_arrive(uint64_t const* smem_ptr) { "cp.async.mbarrier.arrive.shared::cta.b64 [%0];\n\t" "}" : - : "r"(smem_addr)); + : "r"(smem_addr) + : "memory"); cutlass::arch::synclog_emit_cpasync_barrier_arrive(__LINE__, smem_addr); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -752,7 +766,8 @@ void cpasync_barrier_arrive_noinc(uint64_t const* smem_ptr) { "cp.async.mbarrier.arrive.noinc.shared::cta.b64 [%0];\n\t" "}" : - : "r"(smem_addr)); + : "r"(smem_addr) + : "memory"); cutlass::arch::synclog_emit_cpasync_barrier_arrive(__LINE__, smem_addr); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -769,7 +784,8 @@ void umma_arrive(uint64_t const* smem_ptr) { if (cute::elect_one_sync()) { asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];" : - :"r"(bar_intptr)); + :"r"(bar_intptr) + : "memory"); } #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -784,7 +800,8 @@ void umma_arrive_2x1SM(uint64_t const* smem_ptr) { if (cute::elect_one_sync()) { asm volatile("tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.b64 [%0];" : - :"r"(bar_intptr)); + :"r"(bar_intptr) + : "memory"); } #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -802,7 +819,8 @@ void umma_arrive_multicast(uint64_t const* smem_ptr, uint16_t cta_mask) { "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1; \n\t" "}" : - :"r"(bar_intptr), "h"(cta_mask)); + :"r"(bar_intptr), "h"(cta_mask) + : "memory"); } #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -820,7 +838,8 @@ void umma_arrive_multicast_2x1SM(uint64_t const* smem_ptr, uint16_t cta_mask) { "tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1; \n\t" "}" : - :"r"(bar_intptr), "h"(cta_mask)); + :"r"(bar_intptr), "h"(cta_mask) + : "memory"); } #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -840,7 +859,8 @@ void umma_arrive_multicast_no_elect(uint64_t const* smem_ptr, uint16_t cta_mask) "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], lo; \n\t" "}" : - :"r"(bar_intptr), "r"(uint32_t(cta_mask))); + :"r"(bar_intptr), "r"(uint32_t(cta_mask)) + : "memory"); #elif defined(__CUDA_ARCH__) CUTLASS_NOT_IMPLEMENTED(); #endif @@ -859,7 +879,8 @@ void umma_arrive_multicast_2x1SM_no_elect(uint64_t const* smem_ptr, uint16_t cta "tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], lo; \n\t" "}" : - :"r"(bar_intptr), "r"(uint32_t(cta_mask))); + :"r"(bar_intptr), "r"(uint32_t(cta_mask)) + : "memory"); #else CUTLASS_NOT_IMPLEMENTED(); #endif @@ -875,7 +896,8 @@ void umma_arrive_2x1SM_sm0(uint64_t const* smem_ptr) { "mbarrier.arrive.shared::cluster.b64 _, [%0];\n\t" "}" : - : "r"(bar_intptr)); + : "r"(bar_intptr) + : "memory"); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); @@ -888,7 +910,8 @@ CUTE_DEVICE static void fence_view_async_tmem_load() { "{\n\t" "tcgen05.wait::ld.sync.aligned; \n" "}" - ::); + :: + : "memory"); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); #endif @@ -900,7 +923,8 @@ CUTE_DEVICE static void fence_view_async_tmem_store() { "{\n\t" "tcgen05.wait::st.sync.aligned; \n" "}" - ::); + :: + : "memory"); #elif defined(__CUDA_ARCH__) asm volatile ("brkpt;\n" ::); #endif diff --git a/include/cutlass/detail/collective/moe_stride_utils.hpp b/include/cutlass/detail/collective/moe_stride_utils.hpp new file mode 100644 index 00000000..e963ae6e --- /dev/null +++ b/include/cutlass/detail/collective/moe_stride_utils.hpp @@ -0,0 +1,99 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 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. + * + **************************************************************************************************/ +#pragma once + +#include "cutlass/cutlass.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// +namespace cutlass { +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +CUTLASS_HOST_DEVICE +cute::Stride, int64_t> +make_internal_packed_stride(cute::Stride, int64_t> s, cute::Shape shape_MKL) { + static_assert(std::is_integral_v, + "Stride must have an integral type so it can be set dynamically. Static strides not supported."); + auto s_copy = s; + cute::get<0>(s_copy) = static_cast(cute::get<1>(shape_MKL)); + int batch_count = cute::get<2>(shape_MKL); + if (batch_count > 1) { + cute::get<2>(s_copy) = static_cast(cute::get<0>(shape_MKL) * cute::get<1>(shape_MKL)); + } + else { + cute::get<2>(s_copy) = static_cast(0); + } + return s_copy; +} + +template +CUTLASS_HOST_DEVICE +cute::Stride, IntT, int64_t> +make_internal_packed_stride(cute::Stride, IntT, int64_t> s, cute::Shape shape_MKL) { + static_assert(std::is_integral_v, + "Stride must have an integral type so it can be set dynamically. Static strides not supported."); + auto s_copy = s; + cute::get<1>(s_copy) = static_cast(cute::get<0>(shape_MKL)); + int batch_count = cute::get<2>(shape_MKL); + if (batch_count > 1) { + cute::get<2>(s_copy) = static_cast(cute::get<0>(shape_MKL) * cute::get<1>(shape_MKL)); + } + else { + cute::get<2>(s_copy) = static_cast(0); + } + return s_copy; +} + +template +CUTLASS_HOST_DEVICE +cute::Stride, cute::Int<0>> +make_internal_packed_stride(cute::Stride, cute::Int<0>> s, cute::Shape shape_MKL) { + static_assert(std::is_integral_v, + "Stride must have an integral type so it can be set dynamically. Static strides not supported."); + auto s_copy = s; + cute::get<0>(s_copy) = static_cast(cute::get<1>(shape_MKL)); + return s_copy; +} + +template +CUTLASS_HOST_DEVICE +cute::Stride, StrideIntT, cute::Int<0>> +make_internal_packed_stride(cute::Stride, StrideIntT, cute::Int<0>> s, cute::Shape shape_MKL) { + static_assert(std::is_integral_v, + "Stride must have an integral type so it can be set dynamically. Static strides not supported."); + auto s_copy = s; + cute::get<1>(s_copy) = static_cast(cute::get<0>(shape_MKL)); + return s_copy; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// +} +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/epilogue/collective/builders/sm100_builder.inl b/include/cutlass/epilogue/collective/builders/sm100_builder.inl index 7c14ac33..b0490202 100644 --- a/include/cutlass/epilogue/collective/builders/sm100_builder.inl +++ b/include/cutlass/epilogue/collective/builders/sm100_builder.inl @@ -1042,7 +1042,8 @@ sm100_dense_compute_tile_shape_or_override() { constexpr int N_min_D = (detail::is_m_major()) ? 8 * WarpN : (sizeof_bits_v == 6) ? 128 * WarpN // TMA store only supports SW128B for FP6 data type : 128 / sizeof_bits_v * WarpN; - constexpr int N = cute::min(CtaN, cute::max(N_perf, N_min_C, N_min_D)); + constexpr int N_tmp = cute::min(CtaN, cute::max(N_perf, N_min_C, N_min_D)); + constexpr int N = CtaN % N_tmp == 0 ? N_tmp : CtaN; static_assert(CtaN >= N_min_C && CtaN >= N_min_D, "CTA tile too small"); // stride by tmem warp layout and return a by-mode tiler diff --git a/include/cutlass/epilogue/collective/detail.hpp b/include/cutlass/epilogue/collective/detail.hpp index 407ea27d..46557eab 100644 --- a/include/cutlass/epilogue/collective/detail.hpp +++ b/include/cutlass/epilogue/collective/detail.hpp @@ -856,7 +856,7 @@ public: TensorStorage& shared_tensors, TensorMap tensormap ) - { + { auto [acc_state_next] = (*this).template operator()( acc_pipeline, acc_pipe_consumer_state, diff --git a/include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp b/include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp index 852deb7f..a608c510 100644 --- a/include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp +++ b/include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp @@ -45,6 +45,7 @@ #include "cutlass/epilogue/fusion/callbacks.hpp" #include "cutlass/epilogue/fusion/sm100_callbacks_tma_warpspecialized.hpp" #include "cutlass/detail/layout.hpp" +#include "cutlass/detail/collective/moe_stride_utils.hpp" #include "cutlass/trace.h" #include "cute/tensor.hpp" @@ -1473,26 +1474,40 @@ public: if constexpr (IsLoad) { if constexpr (is_source_supported) { + ElementC const* ptr_C = nullptr; + Tensor tensor_c = make_tensor(ptr_C, make_layout(make_shape(M,N,Int<1>{}), InternalStrideC{})); if (params.dC != nullptr) { - ElementC const* ptr_C = nullptr; - Tensor tensor_c = make_tensor(ptr_C, make_layout(make_shape(M,N,Int<1>{}), params.dC[next_group])); - - cute::detail::fill_tma_gmem_shape_stride(params.tma_load_c, tensor_c, - prob_shape, prob_stride); - // Convert strides to byte strides - for (uint64_t& stride : prob_stride) { - stride = (stride * sizeof_bits_v) / 8; - } - cute::tma_descriptor_replace_dims_strides_in_shared_mem(shared_tensormaps.smem_tensormap_C, - prob_shape, - prob_stride); + tensor_c = make_tensor(ptr_C, make_layout(make_shape(M,N,Int<1>{}), params.dC[next_group])); } + else { + auto internal_shape_c = make_shape(static_cast(M), static_cast(N), 1); + InternalStrideC stride_c = make_internal_packed_stride(InternalStrideC{}, internal_shape_c); + tensor_c = make_tensor(ptr_C, make_layout(make_shape(M,N,Int<1>{}), stride_c)); + } + + cute::detail::fill_tma_gmem_shape_stride(params.tma_load_c, tensor_c, + prob_shape, prob_stride); + // Convert strides to byte strides + for (uint64_t& stride : prob_stride) { + stride = (stride * sizeof_bits_v) / 8; + } + cute::tma_descriptor_replace_dims_strides_in_shared_mem(shared_tensormaps.smem_tensormap_C, + prob_shape, + prob_stride); + } } else if constexpr (is_destination_supported) { ElementD const* ptr_D = nullptr; - Tensor tensor_d = make_tensor(ptr_D, make_layout(make_shape(M,N,Int<1>{}), params.dD[next_group])); - + Tensor tensor_d = make_tensor(ptr_D, make_layout(make_shape(M,N,Int<1>{}), InternalStrideD{})); + if (params.dD != nullptr) { + tensor_d = make_tensor(ptr_D, make_layout(make_shape(M,N,Int<1>{}), params.dD[next_group])); + } + else { + auto internal_shape_d = make_shape(static_cast(M), static_cast(N), 1); + InternalStrideD stride_d = make_internal_packed_stride(InternalStrideD{}, internal_shape_d); + tensor_d = make_tensor(ptr_D, make_layout(make_shape(M,N,Int<1>{}), stride_d)); + } cute::detail::fill_tma_gmem_shape_stride(params.tma_store_d, tensor_d, prob_shape, prob_stride); // Convert strides to byte strides @@ -1501,8 +1516,8 @@ public: } cute::tma_descriptor_replace_dims_strides_in_shared_mem(shared_tensormaps.smem_tensormap_D, - prob_shape, - prob_stride); + prob_shape, + prob_stride); } } diff --git a/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp b/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp index 062b9a8b..07cfef69 100644 --- a/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp +++ b/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp @@ -751,13 +751,13 @@ public: ++store_pipe_producer_state; ++issued_stores; - // Wait for the next smem buffer to be available - if (issue_tma_store) { - store_pipeline.producer_acquire(store_pipe_producer_state); - } - synchronize(); - if constexpr (ReuseSmemC) { + + // Wait for the next smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + synchronize(); // producer_acquire returns when at most StagesD-1 committed stores are pending bool store_finished = issued_stores > StorePipeline::UnacquiredStages; // Let dma warp know earliest smem buffer is consumed and empty after StagesD producer commits @@ -866,7 +866,13 @@ public: epi_m_prev = epi_m; epi_n_prev = epi_n; } - + if constexpr (not ReuseSmemC) { + // Wait for the smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + synchronize(); + } // Smem reduction callback entry point using current store buffer for workspace cst_callbacks.reduce(sD_epi(_,_,store_pipe_producer_state.index()), synchronize, epi_m, epi_n, is_last_iteration, tRS_rCompute_frg); @@ -885,7 +891,6 @@ public: for (int i = 0; i < size(tRS_rD_frg); ++i) { tRS_rD_frg(i) = cutlass::NumericArrayConverter{}(tRS_rCompute_frg(i)); } - // Copy tile from register to smem if constexpr (is_destination_supported) { copy(tiled_r2s, tRS_rD, tRS_sD(_,_,_,store_pipe_producer_state.index())); diff --git a/include/cutlass/float8.h b/include/cutlass/float8.h index eab0b35f..18f21e3c 100644 --- a/include/cutlass/float8.h +++ b/include/cutlass/float8.h @@ -1637,46 +1637,46 @@ struct numeric_limits : // CUTLASS_HOST_DEVICE -cutlass::float_e4m3_t operator "" _fe4m3(long double x) { +cutlass::float_e4m3_t operator""_fe4m3(long double x) { return cutlass::float_e4m3_t(float(x)); } CUTLASS_HOST_DEVICE -cutlass::float_e4m3_t operator "" _fe4m3(unsigned long long int x) { +cutlass::float_e4m3_t operator""_fe4m3(unsigned long long int x) { return cutlass::float_e4m3_t(int(x)); } CUTLASS_HOST_DEVICE -cutlass::float_ue4m3_t operator "" _fue4m3(long double x) { +cutlass::float_ue4m3_t operator""_fue4m3(long double x) { return cutlass::float_ue4m3_t(float(x)); } CUTLASS_HOST_DEVICE -cutlass::float_ue4m3_t operator "" _fue4m3(unsigned long long int x) { +cutlass::float_ue4m3_t operator""_fue4m3(unsigned long long int x) { return cutlass::float_ue4m3_t(int(x)); } CUTLASS_HOST_DEVICE -cutlass::float_e5m2_t operator "" _fe5m2(long double x) { +cutlass::float_e5m2_t operator""_fe5m2(long double x) { return cutlass::float_e5m2_t(float(x)); } CUTLASS_HOST_DEVICE -cutlass::float_e5m2_t operator "" _fe5m2(unsigned long long int x) { +cutlass::float_e5m2_t operator""_fe5m2(unsigned long long int x) { return cutlass::float_e5m2_t(int(x)); } CUTLASS_HOST_DEVICE -cutlass::float_ue8m0_t operator "" _fue8m0(long double x) +cutlass::float_ue8m0_t operator""_fue8m0(long double x) { return cutlass::float_ue8m0_t(float(x)); } CUTLASS_HOST_DEVICE -cutlass::float_ue8m0_t operator "" _fue8m0(unsigned long long int x) +cutlass::float_ue8m0_t operator""_fue8m0(unsigned long long int x) { return cutlass::float_ue8m0_t(int(x)); } diff --git a/include/cutlass/float_subbyte.h b/include/cutlass/float_subbyte.h index eefab027..9e1be1db 100644 --- a/include/cutlass/float_subbyte.h +++ b/include/cutlass/float_subbyte.h @@ -760,36 +760,36 @@ struct numeric_limits : public float_ // User-defined literals // CUTLASS_HOST_DEVICE -cutlass::float_e2m1_t operator"" _fe2m1(long double x) +cutlass::float_e2m1_t operator""_fe2m1(long double x) { return cutlass::float_e2m1_t(float(x)); } CUTLASS_HOST_DEVICE -cutlass::float_e2m1_t operator"" _fe2m1(unsigned long long int x) +cutlass::float_e2m1_t operator""_fe2m1(unsigned long long int x) { return cutlass::float_e2m1_t(int(x)); } CUTLASS_HOST_DEVICE -cutlass::float_e2m3_t operator"" _fe2m3(long double x) +cutlass::float_e2m3_t operator""_fe2m3(long double x) { return cutlass::float_e2m3_t(float(x)); } CUTLASS_HOST_DEVICE -cutlass::float_e2m3_t operator"" _fe2m3(unsigned long long int x) +cutlass::float_e2m3_t operator""_fe2m3(unsigned long long int x) { return cutlass::float_e2m3_t(int(x)); } CUTLASS_HOST_DEVICE -cutlass::float_e3m2_t operator"" _fe3m2(long double x) +cutlass::float_e3m2_t operator""_fe3m2(long double x) { return cutlass::float_e3m2_t(float(x)); } CUTLASS_HOST_DEVICE -cutlass::float_e3m2_t operator"" _fe3m2(unsigned long long int x) +cutlass::float_e3m2_t operator""_fe3m2(unsigned long long int x) { return cutlass::float_e3m2_t(int(x)); } diff --git a/include/cutlass/gemm/collective/fp8_accumulation.hpp b/include/cutlass/gemm/collective/fp8_accumulation.hpp index 6ff3a944..fe8a7a2d 100644 --- a/include/cutlass/gemm/collective/fp8_accumulation.hpp +++ b/include/cutlass/gemm/collective/fp8_accumulation.hpp @@ -56,7 +56,6 @@ struct GmmaFP8Accumulation { static_assert(is_rmem::value , "Accumulator tensor must be rmem resident."); private: - TensorAccum& accum_; TensorAccum accum_temp_; uint32_t accum_promotion_interval_; // defines the max num of executed MMAs after which accum should be promoted. @@ -65,8 +64,10 @@ private: uint32_t reset_accum_flag_; // accum needs to be zeroed or not. // promote or `add` the partial accumulators to main accumulator (FADD). + template CUTLASS_DEVICE - void promote_core() { + void promote_core(TensorAccumOrig &accum_) { + CUTE_STATIC_ASSERT_V(size(accum_) == size(accum_temp_)); warpgroup_wait<0>(); CUTLASS_PRAGMA_UNROLL for (int i = 0; i < size(accum_); ++i) { @@ -75,8 +76,10 @@ private: } // `multiply` scale the partial accumulators and `add` to main accumulator (FFMA). + template CUTLASS_DEVICE - void scale_core(ElementAccumulator const &scale) { + void scale_core(TensorAccumOrig &accum_, ElementAccumulator const &scale) { + CUTE_STATIC_ASSERT_V(size(accum_) == size(accum_temp_)); CUTLASS_PRAGMA_UNROLL for (int i = 0; i < size(accum_); ++i) { accum_(i) += accum_temp_(i) * scale; @@ -84,16 +87,17 @@ private: } template < + class TensorAccumOrig, class EngineScale, class LayoutScale> CUTLASS_DEVICE - void scale_core(const cute::Tensor &scale) { + void scale_core(TensorAccumOrig &accum_, const cute::Tensor &scale) { using TensorScale = cute::Tensor; static_assert(is_static::value, "Scale Layout should be static"); static_assert(is_rmem::value , "Scale tensor must be rmem resident."); - static_assert(LayoutAccum{}.shape() == LayoutScale{}.shape(), "Accumulator and scale must have same shape."); + CUTE_STATIC_ASSERT_V(size(accum_) == size(accum_temp_)); CUTLASS_PRAGMA_UNROLL for (int i = 0; i < size(accum_); ++i) { @@ -102,12 +106,13 @@ private: } template < + class TensorAccumOrig, class EngineScaleA, class LayoutScaleA, class EngineScaleB, class LayoutScaleB> CUTLASS_DEVICE - void scale_core(const cute::Tensor &scaleA, const cute::Tensor &scaleB) { + void scale_core(TensorAccumOrig &accum_, const cute::Tensor &scaleA, const cute::Tensor &scaleB) { using TensorScaleA = cute::Tensor; using TensorScaleB = cute::Tensor; @@ -116,8 +121,10 @@ private: static_assert(is_rmem::value, "ScaleA tensor must be rmem resident."); static_assert(is_rmem::value, "ScaleB tensor must be rmem resident."); - static_assert(LayoutAccum{}.shape() == LayoutScaleA{}.shape(), "Accumulator and scaleA must have same shape."); - static_assert(LayoutAccum{}.shape() == LayoutScaleB{}.shape(), "Accumulator and scaleB must have same shape."); + + CUTE_STATIC_ASSERT_V(size(accum_) == size(accum_temp_)); + CUTE_STATIC_ASSERT_V(size(accum_) == size(scaleA)); + CUTE_STATIC_ASSERT_V(size(accum_) == size(scaleB)); CUTLASS_PRAGMA_UNROLL for (int i = 0; i < size(accum_); ++i) { @@ -128,16 +135,15 @@ private: public: CUTLASS_DEVICE GmmaFP8Accumulation( - TensorAccum &accum, + TensorAccum &accum_temp, uint32_t accum_promotion_interval, uint32_t mma_count_per_mainloop_iteration) - : accum_(accum), + : accum_temp_(accum_temp), accum_promotion_interval_(accum_promotion_interval), mma_count_per_mainloop_iteration_(mma_count_per_mainloop_iteration), mma_count_(0), reset_accum_flag_(0) { - accum_temp_ = cute::make_fragment_like(accum); } // @@ -160,21 +166,23 @@ public: // /// promote (add) the results from the MMA accumulators to main accumulator if needed. + template CUTLASS_DEVICE - void promote_if_needed() { + void promote_if_needed(TensorAccumOrig &accum_) { mma_count_ += mma_count_per_mainloop_iteration_; reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0); if (reset_accum_flag_) { - promote_core(); + promote_core(accum_); mma_count_ = 0; } } /// promote (add) the residue results from the MMA accumulators to main accumulator if needed. + template CUTLASS_DEVICE - void promote_residue_if_needed() { + void promote_residue_if_needed(TensorAccumOrig &accum_) { if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) { - promote_core(); + promote_core(accum_); } } @@ -183,95 +191,104 @@ public: // /// scale (multiply_add) the results from the MMA accumulators to main accumulator if needed. + template CUTLASS_DEVICE - void scale_if_needed(ElementAccumulator const &scale) { + void scale_if_needed(TensorAccumOrig &accum_, ElementAccumulator const &scale) { mma_count_ += mma_count_per_mainloop_iteration_; reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0); if (reset_accum_flag_) { - scale_core(scale); + scale_core(accum_, scale); mma_count_ = 0; } } template < + class TensorAccumOrig, class EngineScale, class LayoutScale> CUTLASS_DEVICE - void scale_if_needed(const cute::Tensor &scale) { + void scale_if_needed(TensorAccumOrig &accum_, const cute::Tensor &scale) { mma_count_ += mma_count_per_mainloop_iteration_; reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0); if (reset_accum_flag_) { - scale_core(scale); + scale_core(accum_, scale); mma_count_ = 0; } } template < + class TensorAccumOrig, class EngineScaleA, class LayoutScaleA, class EngineScaleB, class LayoutScaleB> CUTLASS_DEVICE - void scale_if_needed(const cute::Tensor &scaleA, const cute::Tensor &scaleB) { + void scale_if_needed(TensorAccumOrig &accum_, const cute::Tensor &scaleA, const cute::Tensor &scaleB) { mma_count_ += mma_count_per_mainloop_iteration_; reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0); if (reset_accum_flag_) { - scale_core(scaleA, scaleB); + scale_core(accum_, scaleA, scaleB); mma_count_ = 0; } } /// scale (multiply_add) the results from the MMA accumulators to main accumulator without checking the counter. + template CUTLASS_DEVICE - void scale(ElementAccumulator const &scale) { - scale_core(scale); + void scale(TensorAccumOrig &accum_, ElementAccumulator const &scale) { + scale_core(accum_, scale); } template < + class TensorAccumOrig, class EngineScale, class LayoutScale> CUTLASS_DEVICE - void scale(const cute::Tensor &scale) { - scale_core(scale); + void scale(TensorAccumOrig &accum_, const cute::Tensor &scale) { + scale_core(accum_, scale); } template < + class TensorAccumOrig, class EngineScaleA, class LayoutScaleA, class EngineScaleB, class LayoutScaleB> CUTLASS_DEVICE - void scale(const cute::Tensor &scaleA, const cute::Tensor &scaleB) { - scale_core(scaleA, scaleB); + void scale(TensorAccumOrig &accum_, const cute::Tensor &scaleA, const cute::Tensor &scaleB) { + scale_core(accum_, scaleA, scaleB); } /// scale (multiply_add) the residue results from the MMA accumulators to main accumulator if needed. + template CUTLASS_DEVICE - void scale_residue_if_needed(ElementAccumulator const &scale) { + void scale_residue_if_needed(TensorAccumOrig &accum_, ElementAccumulator const &scale) { if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) { - scale_core(scale); + scale_core(accum_, scale); } } template < + class TensorAccumOrig, class EngineScale, class LayoutScale> CUTLASS_DEVICE - void scale_residue_if_needed(const cute::Tensor &scale) { + void scale_residue_if_needed(TensorAccumOrig &accum_, const cute::Tensor &scale) { if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) { - scale_core(scale); + scale_core(accum_, scale); } } template < + class TensorAccumOrig, class EngineScaleA, class LayoutScaleA, class EngineScaleB, class LayoutScaleB> CUTLASS_DEVICE - void scale_residue_if_needed(const cute::Tensor &scaleA, const cute::Tensor &scaleB) { + void scale_residue_if_needed(TensorAccumOrig &accum_, const cute::Tensor &scaleA, const cute::Tensor &scaleB) { if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) { - scale_core(scaleA, scaleB); + scale_core(accum_, scaleA, scaleB); } } }; diff --git a/include/cutlass/gemm/collective/sm100_blockscaled_mma_mixed_tma_cpasync_warpspecialized.hpp b/include/cutlass/gemm/collective/sm100_blockscaled_mma_mixed_tma_cpasync_warpspecialized.hpp index 344de4d3..8bbccdb4 100644 --- a/include/cutlass/gemm/collective/sm100_blockscaled_mma_mixed_tma_cpasync_warpspecialized.hpp +++ b/include/cutlass/gemm/collective/sm100_blockscaled_mma_mixed_tma_cpasync_warpspecialized.hpp @@ -166,8 +166,8 @@ struct CollectiveMma< using ElementBMma = typename TiledMma::ValTypeB; using StrideB = remove_cvref_t(StridePairB{}))>; - static constexpr bool IsRuntimeDataTypeA = cute::is_same_v; - static constexpr bool IsRuntimeDataTypeB = cute::is_same_v; + static constexpr bool IsRuntimeDataTypeA = cute::is_same_v or cute::is_same_v; + static constexpr bool IsRuntimeDataTypeB = cute::is_same_v or cute::is_same_v; static_assert(IsRuntimeDataTypeA == IsRuntimeDataTypeB, "ElementA and ElementB should be both runtime or both static."); @@ -308,13 +308,9 @@ struct CollectiveMma< // Host side kernel arguments struct Arguments { ArrayElementA const* ptr_A{nullptr}; - StrideA dA{}; ArrayElementB const* ptr_B{nullptr}; - StrideB dB{}; ElementSF const* ptr_SFA{nullptr}; - LayoutSFA layout_SFA{}; ElementSF const* ptr_SFB{nullptr}; - LayoutSFB layout_SFB{}; RuntimeDataTypeA runtime_data_type_a{}; RuntimeDataTypeB runtime_data_type_b{}; }; @@ -356,7 +352,6 @@ struct CollectiveMma< TMA_SFB tma_load_sfb; ArrayElementB const* ptr_B{nullptr}; - StrideB dB{}; LayoutSFA layout_SFA; LayoutSFB layout_SFB; @@ -392,9 +387,15 @@ struct CollectiveMma< auto ptr_A = recast_ptr(args.ptr_A); auto ptr_B = recast_ptr(args.ptr_B); - Tensor tensor_a = make_tensor(ptr_A, make_layout(make_shape(M,K,L), args.dA)); - Tensor tensor_sfa = make_tensor(args.ptr_SFA, args.layout_SFA); - Tensor tensor_sfb = make_tensor(args.ptr_SFB, args.layout_SFB); + const auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL); + const auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL); + + auto shape_a = make_shape(M, K, L); + auto stride_a = cutlass::make_internal_packed_stride(StrideA{}, shape_a); + + Tensor tensor_a = make_tensor(ptr_A, make_layout(shape_a, stride_a)); + Tensor tensor_sfa = make_tensor(args.ptr_SFA, layout_SFA); + Tensor tensor_sfb = make_tensor(args.ptr_SFB, layout_SFB); auto cluster_layout_vmnk = tiled_divide(make_layout(ClusterShape{}), make_tile(typename TiledMma::AtomThrID{})); auto cluster_layout_sfb_vmnk = tiled_divide(make_layout(ClusterShape{}), make_tile(typename TiledMma_SF::AtomThrID{})); @@ -426,9 +427,8 @@ struct CollectiveMma< tma_load_sfa, tma_load_sfb, args.ptr_B, - args.dB, - args.layout_SFA, - args.layout_SFB, + layout_SFA, + layout_SFB, args.runtime_data_type_a, args.runtime_data_type_b }; @@ -458,19 +458,6 @@ struct CollectiveMma< CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for CpAsync.\n"); } - // Check for SFA SFB layout requirement - const auto layout_sfa_ref = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL); - const auto layout_sfb_ref = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL); - implementable = implementable && (layout_sfa_ref == args.layout_SFA); - if (!implementable) { - CUTLASS_TRACE_HOST(" CAN IMPLEMENT: layout_SFA mismatch, layout_SFA needs to be K-major\n"); - } - - implementable = implementable && (layout_sfb_ref == args.layout_SFB); - if (!implementable) { - CUTLASS_TRACE_HOST(" CAN IMPLEMENT: layout_SFB mismatch, layout_SFB needs to be K-major\n"); - } - return implementable; } @@ -628,10 +615,14 @@ struct CollectiveMma< // Separate out problem shape for convenience auto [M,N,K,L] = problem_shape_MNKL; + // Setting the stride of B + auto shape_b = make_shape(N, K, L); + StrideB stride_b = cutlass::make_internal_packed_stride(StrideB{}, shape_b); + // convert to subptr iterator if necessary auto ptr_B = recast_ptr(params.ptr_B); // Represent the full tensors - Tensor mB_nkl = make_tensor(make_gmem_ptr(ptr_B), make_shape(N,K,L), params.dB); //(n,k,l) + Tensor mB_nkl = make_tensor(make_gmem_ptr(ptr_B), shape_b, stride_b); //(n,k,l) // Partition for cpasync Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l) @@ -1032,8 +1023,6 @@ protected: RuntimeDataTypeA runtime_data_type_a_{}; RuntimeDataTypeB runtime_data_type_b_{}; - // ClusterShape cluster_shape_; - // uint32_t block_rank_in_cluster_; }; ///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_rcggemm.hpp b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_rcggemm.hpp index 2eb48fb8..0e58264f 100644 --- a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_rcggemm.hpp +++ b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_rcggemm.hpp @@ -50,6 +50,8 @@ #include "cute/algorithm/gemm.hpp" #include "cute/numeric/arithmetic_tuple.hpp" +#include "cutlass/detail/collective/moe_stride_utils.hpp" + ///////////////////////////////////////////////////////////////////////////////////////////////// namespace cutlass::gemm::collective { @@ -257,9 +259,7 @@ struct CollectiveMma< // Host side kernel arguments struct Arguments { ArrayElementA const* ptr_A{nullptr}; - InternalStrideA dA{}; ArrayElementB const** ptr_B{nullptr}; - StrideB dB{}; RuntimeDataTypeA runtime_data_type_a{}; RuntimeDataTypeB runtime_data_type_b{}; }; @@ -296,9 +296,7 @@ struct CollectiveMma< RuntimeDataTypeB runtime_data_type_b; cute::TmaDescriptor* tensormaps; ArrayElementA const* ptr_A; - InternalStrideA dA; ArrayElementB const** ptr_B; - StrideB dB; }; CUTLASS_DEVICE @@ -343,7 +341,9 @@ struct CollectiveMma< init_M = get<0>(problem_shape_MNK); init_K_A = get<2>(problem_shape_MNK); - InternalStrideA stride_a = args.dA; + auto shape_a = make_shape(init_M, init_K_A, problem_shapes.groups()); + InternalStrideA stride_a = cutlass::make_internal_packed_stride(InternalStrideA{}, shape_a); + InternalStrideB stride_b = InternalStrideB{}; // Batches/Groups are managed by using appropriate pointers to input matrices. @@ -397,9 +397,7 @@ struct CollectiveMma< args.runtime_data_type_b, reinterpret_cast(workspace), args.ptr_A, - args.dA, - reinterpret_cast(args.ptr_B), - args.dB + reinterpret_cast(args.ptr_B) }; } @@ -812,7 +810,10 @@ struct CollectiveMma< cute::array prob_stride_B = {0,0,0,0,0}; TmaInternalElementB const* ptr_B = nullptr; - Tensor tensor_b = make_tensor(ptr_B, make_shape(N,K,Int<1>{}), mainloop_params.dB[next_group]); + auto internal_shape_b = make_shape(static_cast(N), static_cast(K), 1); + InternalStrideB stride_b = cutlass::make_internal_packed_stride(InternalStrideB{}, internal_shape_b); + + Tensor tensor_b = make_tensor(ptr_B, make_shape(N,K,Int<1>{}), stride_b); cute::detail::fill_tma_gmem_shape_stride(*observed_tma_load_b_, tensor_b, prob_shape_B, prob_stride_B); diff --git a/include/cutlass/gemm/collective/sm100_mma_mixed_tma_cpasync_warpspecialized.hpp b/include/cutlass/gemm/collective/sm100_mma_mixed_tma_cpasync_warpspecialized.hpp index c31ec335..ccead625 100644 --- a/include/cutlass/gemm/collective/sm100_mma_mixed_tma_cpasync_warpspecialized.hpp +++ b/include/cutlass/gemm/collective/sm100_mma_mixed_tma_cpasync_warpspecialized.hpp @@ -247,9 +247,7 @@ struct CollectiveMma< // Host side kernel arguments struct Arguments { ArrayElementA const* ptr_A{nullptr}; - StrideA dA{}; ArrayElementB const* ptr_B{nullptr}; - StrideB dB{}; RuntimeDataTypeA runtime_data_type_a{}; RuntimeDataTypeB runtime_data_type_b{}; }; @@ -271,7 +269,6 @@ struct CollectiveMma< TMA_A tma_load_a; ArrayElementB const* ptr_B{nullptr}; - StrideB dB{}; RuntimeDataTypeA runtime_data_type_a; RuntimeDataTypeB runtime_data_type_b; @@ -299,7 +296,9 @@ struct CollectiveMma< auto ptr_A = recast_ptr(args.ptr_A); auto ptr_B = recast_ptr(args.ptr_B); - Tensor tensor_a = make_tensor(ptr_A, make_layout(make_shape(M,K,L), args.dA)); + auto shape_a = make_shape(M, K, L); + StrideA stride_a = cutlass::make_internal_packed_stride(StrideA{}, shape_a); + Tensor tensor_a = make_tensor(ptr_A, make_layout(shape_a, stride_a)); auto cluster_layout_vmnk = tiled_divide(make_layout(ClusterShape{}), make_tile(typename TiledMma::AtomThrID{})); @@ -314,7 +313,6 @@ struct CollectiveMma< return { tma_load_a, args.ptr_B, - args.dB, args.runtime_data_type_a, args.runtime_data_type_b }; @@ -443,7 +441,11 @@ struct CollectiveMma< auto [M,N,K,L] = problem_shape_MNKL; // Represent the full tensors - Tensor mB_nkl = make_tensor(make_gmem_ptr(params.ptr_B), make_shape(N,K,L), params.dB); //(n,k,l) + + auto shape_b = make_shape(N, K, L); + StrideB stride_b = cutlass::make_internal_packed_stride(StrideB{}, shape_b); + + Tensor mB_nkl = make_tensor(make_gmem_ptr(params.ptr_B), shape_b, stride_b); //(n,k,l) // Partition for cpasync Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l) @@ -571,14 +573,6 @@ struct CollectiveMma< ProblemShape_MNKL effective_shape ) { - // Unpack from load_inputs - // GTensorB tBgB_nkl = get<0>(load_inputs); - // CTensorB cgB_nk = get<1>(load_inputs); - // STensorB sB = get<2>(load_inputs); - // ProblemShape_MNKL problem_shape_MNKL = get<3>(load_inputs); - // TiledCopyB gmem_to_smem_b_tiled_copy = get<4>(load_inputs); - // ThreadCopyB thr_copy_b = get<5>(load_inputs); - auto [ tBgB_nkl, cgB_nk, sB, // problem_shape_MNKL, diff --git a/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8.hpp b/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8.hpp index 916c6db8..52ef5777 100644 --- a/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8.hpp +++ b/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8.hpp @@ -531,8 +531,8 @@ struct CollectiveMma< int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count); tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; - - GmmaFP8Accumulation accumulation(accum, mainloop_params.mma_promotion_interval, size<2>(tCrA)); + auto accm_temp = cute::make_fragment_like(accum); + GmmaFP8Accumulation accumulation(accm_temp, mainloop_params.mma_promotion_interval, size<2>(tCrA)); warpgroup_fence_operand(accumulation()); CUTLASS_PRAGMA_UNROLL for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0; --k_tile_prologue) @@ -556,7 +556,7 @@ struct CollectiveMma< } warpgroup_commit_batch(); - accumulation.promote_if_needed(); + accumulation.promote_if_needed(accum); ++smem_pipe_read; } @@ -597,7 +597,7 @@ struct CollectiveMma< warpgroup_wait(); warpgroup_fence_operand(accumulation()); - accumulation.promote_if_needed(); + accumulation.promote_if_needed(accum); pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it @@ -606,7 +606,7 @@ struct CollectiveMma< ++smem_pipe_release; } - accumulation.promote_residue_if_needed(); + accumulation.promote_residue_if_needed(accum); warpgroup_fence_operand(accumulation()); } diff --git a/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp b/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp index b6e662be..33aa42be 100644 --- a/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp +++ b/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp @@ -212,6 +212,8 @@ struct CollectiveMma< static_assert(cute::is_same_v, "ElementAccumulator and ElementBlockScale should be same datatype"); + using NumSplitsM = cute::C(TileShape_{}) / 128>; + static_assert(NumSplitsM{} == 1 || NumSplitsM{} == 2); struct SharedStorage { struct TensorStorage : cute::aligned_struct<128, _0> { @@ -687,35 +689,37 @@ struct CollectiveMma< template< + class AccumSlice, class EngineAccum, class LayoutAccum, class ScaleFactor > CUTLASS_DEVICE - void scale_if_needed(GmmaFP8Accumulation& accumulation, ScaleFactor scaleFactor) { + void scale_if_needed(AccumSlice & accum, GmmaFP8Accumulation& accumulation, ScaleFactor scaleFactor) { if constexpr (ScalePromotionInterval != 4) { - accumulation.scale_if_needed(scaleFactor); + accumulation.scale_if_needed(accum, scaleFactor); } else { // avoid unnecessary tests when granularity is the finnest - accumulation.scale(scaleFactor); + accumulation.scale(accum, scaleFactor); } } template< + class AccumSlice, class EngineAccum, class LayoutAccum, class ScaleFactor1, class ScaleFactor2 > CUTLASS_DEVICE - void scale_if_needed(GmmaFP8Accumulation& accumulation, ScaleFactor1 scaleFactor1, ScaleFactor2 scaleFactor2) { + void scale_if_needed(AccumSlice & accum, GmmaFP8Accumulation& accumulation, ScaleFactor1 scaleFactor1, ScaleFactor2 scaleFactor2) { if constexpr (ScalePromotionInterval != 4) { - accumulation.scale_if_needed(scaleFactor1, scaleFactor2); + accumulation.scale_if_needed(accum, scaleFactor1, scaleFactor2); } else { // avoid unnecessary tests when granularity is the finnest - accumulation.scale(scaleFactor1, scaleFactor2); + accumulation.scale(accum, scaleFactor1, scaleFactor2); } } @@ -821,75 +825,26 @@ struct CollectiveMma< // Prologue GMMAs tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; + // Tile accum + + using NumSplitsM_Scale = cute::conditional_t; + static constexpr int ScaleMsPerWave = ScaleMsPerTile == 1 ? 1 : ScaleMsPerTile / NumSplitsM{}; + + auto accum_tiled = tiled_divide(accum, cute::tuple<_1, NumSplitsM>{}); + auto tCrA_tiled = tiled_divide(tCrA, cute::tuple<_1, NumSplitsM>{}); + auto tCsSFA_tiled = tiled_divide(tCsSFA, cute::tuple<_1, NumSplitsM_Scale>{}); + auto tCrSFA_tiled = tiled_divide(tCrSFA, cute::tuple<_1, NumSplitsM_Scale>{}); + auto tCrSFB_tiled = tiled_divide(tCrSFB, cute::tuple<_1, NumSplitsM_Scale>{}); + // Temporary accumulator used by MMA + // On promotion, accumulated values are scaled and copied into `accum` + auto accum_temp = cute::make_fragment_like(accum_tiled(_0{}, _, _, _)); auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read); - GmmaFP8Accumulation accumulation(accum, ScalePromotionInterval, size<2>(tCrA)); - warpgroup_fence_operand(accumulation()); - - if (k_tile_count > 0) { - // WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value) - pipeline.consumer_wait(smem_pipe_read, barrier_token); - - int read_stage = smem_pipe_read.index(); - // Load per block scale values from shared memory to registers - copy(tCsSFA(_,_,_,make_coord(_0{},read_stage)), tCrSFA); - copy(tCsSFB(_,_,_,make_coord(_0{},read_stage)), tCrSFB); - - warpgroup_fence_operand(accumulation()); - warpgroup_arrive(); - // Unroll the K mode manually to set scale D to 1 - CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { - // (V,M) x (V,N) => (V,M,N) - cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation()); - tiled_mma.accumulate_ = GMMA::ScaleOut::One; - } - warpgroup_commit_batch(); - warpgroup_fence_operand(accumulation()); - - - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { - tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{}); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { - ElementBlockScale scale_b = tCrSFB(_0{}); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) { - filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b; - } - } - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { - ElementBlockScale scale_a = tCrSFA(_0{}); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) { - filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a; - } - } - - warpgroup_wait<0>(); - ++smem_pipe_read; - barrier_token = pipeline.consumer_try_wait(smem_pipe_read); - - // Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB` - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { - ElementBlockScale scale_ab = tCrSFA(_0{}); - scale_if_needed(accumulation, scale_ab); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { - scale_if_needed(accumulation, tCrSFA); - } - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { - scale_if_needed(accumulation, tCrSFB); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) { - scale_if_needed(accumulation, tCrSFA, tCrSFB); - } - } + // Secondary accumulator for FP32 accum + GmmaFP8Accumulation accumulation(accum_temp, ScalePromotionInterval, size<2>(tCrA)); warpgroup_fence_operand(accumulation()); - // Mainloop GMMAs - k_tile_count--; CUTLASS_PRAGMA_NO_UNROLL for ( ; k_tile_count > 1; --k_tile_count) @@ -903,73 +858,86 @@ struct CollectiveMma< int read_stage = smem_pipe_read.index(); // Load per block scale values from shared memory to registers (at most twice per block along M and/or N) - copy(tCsSFA(_,_,_,make_coord(_0{}, read_stage)), tCrSFA); - copy(tCsSFB(_,_,_,make_coord(_0{}, read_stage)), tCrSFB); + copy(tCsSFB(_,_,_,make_coord(_0{}, read_stage)), tCrSFB); - if constexpr (ScalePromotionInterval != 4) { - if (accumulation.prepare_if_needed()) { + CUTLASS_PRAGMA_UNROLL + for (int m_split = 0; m_split < NumSplitsM{}; ++m_split) { + auto tCrA_local = tCrA_tiled(m_split, _, _, _, _); + auto tCrSFA_local = tCrSFA_tiled(m_split, _, _, _); + auto tCrSFB_local = tCrSFB_tiled(m_split, _, _, _); + auto accum_local = accum_tiled(m_split, _, _, _); + copy(tCsSFA_tiled(m_split, _, _, _, make_coord(_0{}, read_stage)), tCrSFA_local); + bool is_last = (m_split == NumSplitsM{} - 1); + + if constexpr (ScalePromotionInterval != 4) { + if (accumulation.prepare_if_needed()) { + tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; + } + } + else { + // Always zero out the accumulator for finest granularity tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; } - } - else { - // Always zero out the accumulator for finest granularity - tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; - } - warpgroup_fence_operand(accumulation()); - warpgroup_arrive(); - // Unroll the K mode manually to set scale D to 1 - CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { - // (V,M) x (V,N) => (V,M,N) - cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation()); - tiled_mma.accumulate_ = GMMA::ScaleOut::One; - } - warpgroup_commit_batch(); - - /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed - warpgroup_fence_operand(accumulation()); - - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { - tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{}); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { - ElementBlockScale scale_b = tCrSFB(_0{}); + warpgroup_fence_operand(accumulation()); + warpgroup_arrive(); + // Unroll the K mode manually to set scale D to 1 CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) { - filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b; + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + // (V,M) x (V,N) => (V,M,N) + cute::gemm(tiled_mma, tCrA_local(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation()); + tiled_mma.accumulate_ = GMMA::ScaleOut::One; } - } - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { - ElementBlockScale scale_a = tCrSFA(_0{}); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) { - filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a; + warpgroup_commit_batch(); + + /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed + warpgroup_fence_operand(accumulation()); + + + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) { + tCrSFA_local(_0{}) = tCrSFA_local(_0{}) * tCrSFB(_0{}); + } + if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) { + ElementBlockScale scale_b = tCrSFB(_0{}); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(filter_zeros(tCrSFA_local)); i++) { + filter_zeros(tCrSFA_local)(i) = filter_zeros(tCrSFA_local)(i) * scale_b; + } + } + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) { + ElementBlockScale scale_a = tCrSFA_local(_0{}); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(filter_zeros(tCrSFB_local)); i++) { + filter_zeros(tCrSFB_local)(i) = filter_zeros(tCrSFB_local)(i) * scale_a; + } } - } - warpgroup_wait<0>(); - pipeline.consumer_release(smem_pipe_release); // Unlock previous tile - ++smem_pipe_read; - barrier_token = pipeline.consumer_try_wait(smem_pipe_read); + warpgroup_wait<0>(); + if (is_last) { + pipeline.consumer_release(smem_pipe_release); // Unlock previous tile + ++smem_pipe_read; + barrier_token = pipeline.consumer_try_wait(smem_pipe_read); + } + // Block scale the accumulators with reg tensor `tCrSFA_local` and `tCrSFB` + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) { + ElementBlockScale scale_ab = tCrSFA_local(_0{}); + scale_if_needed(accum_local, accumulation, scale_ab); + } + if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) { + scale_if_needed(accum_local, accumulation, tCrSFA_local); + } + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) { + scale_if_needed(accum_local, accumulation, tCrSFB_local); + } + if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile > 1) { + scale_if_needed(accum_local, accumulation, tCrSFA_local, tCrSFB_local); + } - // Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB` - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { - ElementBlockScale scale_ab = tCrSFA(_0{}); - scale_if_needed(accumulation, scale_ab); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { - scale_if_needed(accumulation, tCrSFA); - } - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { - scale_if_needed(accumulation, tCrSFB); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) { - scale_if_needed(accumulation, tCrSFA, tCrSFB); - } - - // Advance smem_pipe_read and smem_pipe_release - ++smem_pipe_release; + if (is_last) { + // Advance smem_pipe_read and smem_pipe_release + ++smem_pipe_release; + } + } // end for (m_split) } if (k_tile_count > 0) { @@ -981,97 +949,101 @@ struct CollectiveMma< int read_stage = smem_pipe_read.index(); // Load per block scale values from shared memory to registers (at most twice per block along M and/or N) - copy(tCsSFA(_,_,_,make_coord(_0{}, read_stage)), tCrSFA); copy(tCsSFB(_,_,_,make_coord(_0{}, read_stage)), tCrSFB); + CUTLASS_PRAGMA_UNROLL + for (int m_split = 0; m_split < NumSplitsM{}; ++m_split) { + auto tCrA_local = tCrA_tiled(m_split, _, _, _, _); + auto tCrSFA_local = tCrSFA_tiled(m_split, _, _, _); + auto tCrSFB_local = tCrSFB_tiled(m_split, _, _, _); + auto accum_local = accum_tiled(m_split, _, _, _); + copy(tCsSFA_tiled(m_split, _, _, _, make_coord(_0{}, read_stage)), tCrSFA_local); + bool is_last = (m_split == NumSplitsM{} - 1); - if constexpr (ScalePromotionInterval != 4) { - if (accumulation.prepare_if_needed()) { + if constexpr (ScalePromotionInterval != 4) { + if (accumulation.prepare_if_needed()) { + tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; + } + } + else { + // Always zero out the accumulator for finest granularity tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; } - } - else { - // Always zero out the accumulator for finest granularity - tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; - } - warpgroup_fence_operand(accumulation()); - warpgroup_arrive(); - // Unroll the K mode manually to set scale D to 1 - CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { - // (V,M) x (V,N) => (V,M,N) - cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation()); - tiled_mma.accumulate_ = GMMA::ScaleOut::One; - } - warpgroup_commit_batch(); - - /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed - warpgroup_fence_operand(accumulation()); - - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { - tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{}); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { - ElementBlockScale scale_b = tCrSFB(_0{}); + warpgroup_fence_operand(accumulation()); + warpgroup_arrive(); + // Unroll the K mode manually to set scale D to 1 CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) { - filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b; + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + // (V,M) x (V,N) => (V,M,N) + cute::gemm(tiled_mma, tCrA_local(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation()); + tiled_mma.accumulate_ = GMMA::ScaleOut::One; } - } - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { - ElementBlockScale scale_a = tCrSFA(_0{}); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) { - filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a; + warpgroup_commit_batch(); + + /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed + warpgroup_fence_operand(accumulation()); + + + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) { + tCrSFA_local(_0{}) = tCrSFA_local(_0{}) * tCrSFB(_0{}); + } + if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) { + ElementBlockScale scale_b = tCrSFB(_0{}); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(filter_zeros(tCrSFA_local)); i++) { + filter_zeros(tCrSFA_local)(i) = filter_zeros(tCrSFA_local)(i) * scale_b; + } + } + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) { + ElementBlockScale scale_a = tCrSFA_local(_0{}); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(filter_zeros(tCrSFB_local)); i++) { + filter_zeros(tCrSFB_local)(i) = filter_zeros(tCrSFB_local)(i) * scale_a; + } + } + warpgroup_wait<0>(); + if (is_last) { + pipeline.consumer_release(smem_pipe_release); // Unlock previous tile + } + // Block scale the accumulators with reg tensor `tCrSFA_local` and `tCrSFB` + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) { + ElementBlockScale scale_ab = tCrSFA_local(_0{}); + scale_if_needed(accum_local, accumulation, scale_ab); + } + if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) { + scale_if_needed(accum_local, accumulation, tCrSFA_local); + } + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) { + scale_if_needed(accum_local, accumulation, tCrSFB_local); + } + if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile > 1) { + scale_if_needed(accum_local, accumulation, tCrSFA_local, tCrSFB_local); + } + if constexpr (ScalePromotionInterval != 4) { + // residues only exists when granularity is not the finnest + if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { + ElementBlockScale scale_ab = tCrSFA_local(_0{}); + accumulation.scale_residue_if_needed(accum_local, scale_ab); + } + if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { + accumulation.scale_residue_if_needed(accum_local, tCrSFA_local); + } + if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { + accumulation.scale_residue_if_needed(accum_local, tCrSFB_local); + } + if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) { + accumulation.scale_residue_if_needed(accum_local, tCrSFA_local, tCrSFB_local); + } } - } - warpgroup_wait<0>(); - pipeline.consumer_release(smem_pipe_release); // Unlock previous tile - // Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB` - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { - ElementBlockScale scale_ab = tCrSFA(_0{}); - scale_if_needed(accumulation, scale_ab); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { - scale_if_needed(accumulation, tCrSFA); - } - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { - scale_if_needed(accumulation, tCrSFB); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) { - scale_if_needed(accumulation, tCrSFA, tCrSFB); - } + warpgroup_fence_operand(accumulation()); + } // end for (m_split) } - if constexpr (ScalePromotionInterval != 4) { - // residues only exists when granularity is not the finnest - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { - ElementBlockScale scale_ab = tCrSFA(_0{}); - accumulation.scale_residue_if_needed(scale_ab); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { - accumulation.scale_residue_if_needed(tCrSFA); - } - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { - accumulation.scale_residue_if_needed(tCrSFB); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) { - accumulation.scale_residue_if_needed(tCrSFA, tCrSFB); - } - } - - warpgroup_fence_operand(accumulation()); - } /// Perform a Consumer Epilogue to release all buffers CUTLASS_DEVICE void mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) { - if (k_tile_count > 0) { - // The pipeline is not released in the first iteration - smem_pipe_release.advance(k_tile_count - 1); - pipeline.consumer_release(smem_pipe_release); - } } // diff --git a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp index c7ea65a6..1f2089e8 100644 --- a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp +++ b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp @@ -481,8 +481,8 @@ struct CollectiveMma< int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count); tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; - - GmmaFP8Accumulation accumulation(accum, mainloop_params.mma_promotion_interval, size<2>(tCrA)); + auto accm_temp = cute::make_fragment_like(accum); + GmmaFP8Accumulation accumulation(accm_temp, mainloop_params.mma_promotion_interval, size<2>(tCrA)); warpgroup_fence_operand(accumulation()); CUTLASS_PRAGMA_UNROLL for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0; --k_tile_prologue) @@ -506,7 +506,7 @@ struct CollectiveMma< } warpgroup_commit_batch(); - accumulation.promote_if_needed(); + accumulation.promote_if_needed(accum); ++smem_pipe_read; } @@ -547,7 +547,7 @@ struct CollectiveMma< warpgroup_wait(); warpgroup_fence_operand(accumulation()); - accumulation.promote_if_needed(); + accumulation.promote_if_needed(accum); pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it @@ -556,7 +556,7 @@ struct CollectiveMma< ++smem_pipe_release; } - accumulation.promote_residue_if_needed(); + accumulation.promote_residue_if_needed(accum); warpgroup_fence_operand(accumulation()); } diff --git a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp index 48ddf7a0..c327c1f6 100644 --- a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp +++ b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp @@ -33,7 +33,9 @@ #include "cutlass/cutlass.h" #include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/collective/fp8_accumulation.hpp" #include "cutlass/trace.h" +#include "cutlass/pipeline/pipeline.hpp" #include "cutlass/numeric_types.h" #include "cute/arch/cluster_sm90.hpp" @@ -203,6 +205,9 @@ struct CollectiveMma< static_assert(cute::is_same_v, "ElementAccumulator and ElementBlockScale should be same datatype"); + using NumSplitsM = cute::C(TileShape_{}) / 128>; + static_assert(NumSplitsM{} == 1 || NumSplitsM{} == 2); + struct SharedStorage { struct TensorStorage : cute::aligned_struct<128> { @@ -720,34 +725,36 @@ struct CollectiveMma< } template< + class AccumSlice, class EngineAccum, class LayoutAccum, class ScaleFactor > CUTLASS_DEVICE - void scale_if_needed(GmmaFP8Accumulation& accumulation, ScaleFactor scaleFactor) { + void scale_if_needed(AccumSlice & accum, GmmaFP8Accumulation& accumulation, ScaleFactor scaleFactor) { if constexpr (ScalePromotionInterval != 4) { - accumulation.scale_if_needed(scaleFactor); + accumulation.scale_if_needed(accum, scaleFactor); } else { // avoid unnecessary tests when granularity is the finnest - accumulation.scale(scaleFactor); + accumulation.scale(accum, scaleFactor); } } template< + class AccumSlice, class EngineAccum, class LayoutAccum, class ScaleFactor1, class ScaleFactor2 > CUTLASS_DEVICE - void scale_if_needed(GmmaFP8Accumulation& accumulation, ScaleFactor1 scaleFactor1, ScaleFactor2 scaleFactor2) { + void scale_if_needed(AccumSlice & accum, GmmaFP8Accumulation& accumulation, ScaleFactor1 scaleFactor1, ScaleFactor2 scaleFactor2) { if constexpr (ScalePromotionInterval != 4) { - accumulation.scale_if_needed(scaleFactor1, scaleFactor2); + accumulation.scale_if_needed(accum, scaleFactor1, scaleFactor2); } else { // avoid unnecessary tests when granularity is the finnest - accumulation.scale(scaleFactor1, scaleFactor2); + accumulation.scale(accum, scaleFactor1, scaleFactor2); } } @@ -852,70 +859,27 @@ struct CollectiveMma< // Prologue GMMAs tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; + // Tile accum + using NumSplitsM_Scale = cute::conditional_t; + static constexpr int ScaleMsPerWave = ScaleMsPerTile == 1 ? 1 : ScaleMsPerTile / NumSplitsM{}; + + auto accum_tiled = tiled_divide(accum, cute::tuple<_1, NumSplitsM>{}); + auto tCrA_tiled = tiled_divide(tCrA, cute::tuple<_1, NumSplitsM>{}); + auto tCsSFA_tiled = tiled_divide(tCsSFA, cute::tuple<_1, NumSplitsM_Scale>{}); + auto tCrSFA_tiled = tiled_divide(tCrSFA, cute::tuple<_1, NumSplitsM_Scale>{}); + auto tCrSFB_tiled = tiled_divide(tCrSFB, cute::tuple<_1, NumSplitsM_Scale>{}); + // Temporary accumulator used by MMA + // On promotion, accumulated values are scaled and copied into `accum` + auto accum_temp = cute::make_fragment_like(accum_tiled(_0{}, _, _, _)); // WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value) auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read); pipeline.consumer_wait(smem_pipe_read, barrier_token); - GmmaFP8Accumulation accumulation(accum, ScalePromotionInterval, size<2>(tCrA)); - warpgroup_fence_operand(accumulation()); - { - int read_stage = smem_pipe_read.index(); - - // Load per block scale values from shared memory to registers - copy(tCsSFA(_,_,_,make_coord(_0{}, read_stage)), tCrSFA); - copy(tCsSFB(_,_,_,make_coord(_0{}, read_stage)), tCrSFB); - - warpgroup_fence_operand(accumulation()); - warpgroup_arrive(); - // Unroll the K mode manually to set scale D to 1 - CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { - // (V,M) x (V,N) => (V,M,N) - cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation()); - tiled_mma.accumulate_ = GMMA::ScaleOut::One; - } - warpgroup_commit_batch(); - warpgroup_fence_operand(accumulation()); - - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { - tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{}); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { - ElementBlockScale scale_b = tCrSFB(_0{}); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) { - filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b; - } - } - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { - ElementBlockScale scale_a = tCrSFA(_0{}); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) { - filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a; - } - } - warpgroup_wait<0>(); - ++smem_pipe_read; - barrier_token = pipeline.consumer_try_wait(smem_pipe_read); - // Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB` - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { - ElementBlockScale scale_ab = tCrSFA(_0{}); - scale_if_needed(accumulation, scale_ab); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { - scale_if_needed(accumulation, tCrSFA); - } - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { - scale_if_needed(accumulation, tCrSFB); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) { - scale_if_needed(accumulation, tCrSFA, tCrSFB); - } - } + // Secondary accumulator for FP32 accum + GmmaFP8Accumulation accumulation(accum_temp, ScalePromotionInterval, size<2>(tCrA)); warpgroup_fence_operand(accumulation()); // Mainloop GMMAs - k_tile_count -= 1; CUTLASS_PRAGMA_NO_UNROLL for ( ; k_tile_count > 1; --k_tile_count) @@ -929,71 +893,86 @@ struct CollectiveMma< int read_stage = smem_pipe_read.index(); // Load per block scale values from shared memory to registers (at most twice per block along M and/or N) - copy(tCsSFA(_,_,_,make_coord(_0{}, read_stage)), tCrSFA); copy(tCsSFB(_,_,_,make_coord(_0{}, read_stage)), tCrSFB); - if constexpr (ScalePromotionInterval != 4) { - if (accumulation.prepare_if_needed()) { + CUTLASS_PRAGMA_UNROLL + for (int m_split = 0; m_split < NumSplitsM{}; ++m_split) { + auto tCrA_local = tCrA_tiled(m_split, _, _, _, _); + auto tCrSFA_local = tCrSFA_tiled(m_split, _, _, _); + auto tCrSFB_local = tCrSFB_tiled(m_split, _, _, _); + auto accum_local = accum_tiled(m_split, _, _, _); + copy(tCsSFA_tiled(m_split, _, _, _, make_coord(_0{}, read_stage)), tCrSFA_local); + bool is_last = (m_split == NumSplitsM{} - 1); + + if constexpr (ScalePromotionInterval != 4) { + if (accumulation.prepare_if_needed()) { + tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; + } + } + else { + // Always zero out the accumulator for finest granularity tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; } - } - else { - // Always zero out the accumulator for finest granularity - tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; - } - warpgroup_fence_operand(accumulation()); - warpgroup_arrive(); - // Unroll the K mode manually to set scale D to 1 - CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { - // (V,M) x (V,N) => (V,M,N) - cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation()); - tiled_mma.accumulate_ = GMMA::ScaleOut::One; - } - warpgroup_commit_batch(); - - /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed - warpgroup_fence_operand(accumulation()); - - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { - tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{}); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { - ElementBlockScale scale_b = tCrSFB(_0{}); + warpgroup_fence_operand(accumulation()); + warpgroup_arrive(); + // Unroll the K mode manually to set scale D to 1 CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) { - filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b; + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + // (V,M) x (V,N) => (V,M,N) + cute::gemm(tiled_mma, tCrA_local(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation()); + tiled_mma.accumulate_ = GMMA::ScaleOut::One; } - } - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { - ElementBlockScale scale_a = tCrSFA(_0{}); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) { - filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a; - } - } - warpgroup_wait<0>(); - pipeline.consumer_release(smem_pipe_release); // Unlock previous tile - ++smem_pipe_read; - barrier_token = pipeline.consumer_try_wait(smem_pipe_read); - // Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB` - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { - ElementBlockScale scale_ab = tCrSFA(_0{}); - scale_if_needed(accumulation, scale_ab); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { - scale_if_needed(accumulation, tCrSFA); - } - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { - scale_if_needed(accumulation, tCrSFB); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) { - scale_if_needed(accumulation, tCrSFA, tCrSFB); - } + warpgroup_commit_batch(); - // Advance smem_pipe_read and smem_pipe_release - ++smem_pipe_release; + /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed + warpgroup_fence_operand(accumulation()); + + + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) { + tCrSFA_local(_0{}) = tCrSFA_local(_0{}) * tCrSFB(_0{}); + } + if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) { + ElementBlockScale scale_b = tCrSFB(_0{}); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(filter_zeros(tCrSFA_local)); i++) { + filter_zeros(tCrSFA_local)(i) = filter_zeros(tCrSFA_local)(i) * scale_b; + } + } + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) { + ElementBlockScale scale_a = tCrSFA_local(_0{}); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(filter_zeros(tCrSFB_local)); i++) { + filter_zeros(tCrSFB_local)(i) = filter_zeros(tCrSFB_local)(i) * scale_a; + } + } + + warpgroup_wait<0>(); + if (is_last) { + pipeline.consumer_release(smem_pipe_release); // Unlock previous tile + ++smem_pipe_read; + barrier_token = pipeline.consumer_try_wait(smem_pipe_read); + } + // Block scale the accumulators with reg tensor `tCrSFA_local` and `tCrSFB` + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) { + ElementBlockScale scale_ab = tCrSFA_local(_0{}); + scale_if_needed(accum_local, accumulation, scale_ab); + } + if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) { + scale_if_needed(accum_local, accumulation, tCrSFA_local); + } + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) { + scale_if_needed(accum_local, accumulation, tCrSFB_local); + } + if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile > 1) { + scale_if_needed(accum_local, accumulation, tCrSFA_local, tCrSFB_local); + } + + if (is_last) { + // Advance smem_pipe_read and smem_pipe_release + ++smem_pipe_release; + } + } // end for (m_split) } if (k_tile_count) { pipeline.consumer_wait(smem_pipe_read, barrier_token); @@ -1005,93 +984,101 @@ struct CollectiveMma< int read_stage = smem_pipe_read.index(); // Load per block scale values from shared memory to registers (at most twice per block along M and/or N) - copy(tCsSFA(_,_,_,make_coord(_0{}, read_stage)), tCrSFA); copy(tCsSFB(_,_,_,make_coord(_0{}, read_stage)), tCrSFB); + CUTLASS_PRAGMA_UNROLL + for (int m_split = 0; m_split < NumSplitsM{}; ++m_split) { + auto tCrA_local = tCrA_tiled(m_split, _, _, _, _); + auto tCrSFA_local = tCrSFA_tiled(m_split, _, _, _); + auto tCrSFB_local = tCrSFB_tiled(m_split, _, _, _); + auto accum_local = accum_tiled(m_split, _, _, _); + copy(tCsSFA_tiled(m_split, _, _, _, make_coord(_0{}, read_stage)), tCrSFA_local); + bool is_last = (m_split == NumSplitsM{} - 1); - if constexpr (ScalePromotionInterval != 4) { - if (accumulation.prepare_if_needed()) { + if constexpr (ScalePromotionInterval != 4) { + if (accumulation.prepare_if_needed()) { + tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; + } + } + else { + // Always zero out the accumulator for finest granularity tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; } - } - else { - // Always zero out the accumulator for finest granularity - tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; - } - warpgroup_fence_operand(accumulation()); - warpgroup_arrive(); - // Unroll the K mode manually to set scale D to 1 - CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { - // (V,M) x (V,N) => (V,M,N) - cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation()); - tiled_mma.accumulate_ = GMMA::ScaleOut::One; - } - warpgroup_commit_batch(); - - /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed - warpgroup_fence_operand(accumulation()); - - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { - tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{}); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { - ElementBlockScale scale_b = tCrSFB(_0{}); + warpgroup_fence_operand(accumulation()); + warpgroup_arrive(); + // Unroll the K mode manually to set scale D to 1 CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) { - filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b; + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + // (V,M) x (V,N) => (V,M,N) + cute::gemm(tiled_mma, tCrA_local(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation()); + tiled_mma.accumulate_ = GMMA::ScaleOut::One; } - } - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { - ElementBlockScale scale_a = tCrSFA(_0{}); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) { - filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a; - } - } - warpgroup_wait<0>(); - pipeline.consumer_release(smem_pipe_release); // Unlock previous tile - // Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB` - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { - ElementBlockScale scale_ab = tCrSFA(_0{}); - scale_if_needed(accumulation, scale_ab); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { - scale_if_needed(accumulation, tCrSFA); - } - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { - scale_if_needed(accumulation, tCrSFB); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) { - scale_if_needed(accumulation, tCrSFA, tCrSFB); - } - } - if constexpr (ScalePromotionInterval != 4) { - // residues only exists when granularity is not the finnest - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { - ElementBlockScale scale_ab = tCrSFA(_0{}); - accumulation.scale_residue_if_needed(scale_ab); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { - accumulation.scale_residue_if_needed(tCrSFA); - } - if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { - accumulation.scale_residue_if_needed(tCrSFB); - } - if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) { - accumulation.scale_residue_if_needed(tCrSFA, tCrSFB); - } - } + warpgroup_commit_batch(); - warpgroup_fence_operand(accumulation()); + /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed + warpgroup_fence_operand(accumulation()); + + + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) { + tCrSFA_local(_0{}) = tCrSFA_local(_0{}) * tCrSFB(_0{}); + } + if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) { + ElementBlockScale scale_b = tCrSFB(_0{}); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(filter_zeros(tCrSFA_local)); i++) { + filter_zeros(tCrSFA_local)(i) = filter_zeros(tCrSFA_local)(i) * scale_b; + } + } + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) { + ElementBlockScale scale_a = tCrSFA_local(_0{}); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(filter_zeros(tCrSFB_local)); i++) { + filter_zeros(tCrSFB_local)(i) = filter_zeros(tCrSFB_local)(i) * scale_a; + } + } + warpgroup_wait<0>(); + if (is_last) { + pipeline.consumer_release(smem_pipe_release); // Unlock previous tile + } + // Block scale the accumulators with reg tensor `tCrSFA_local` and `tCrSFB` + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) { + ElementBlockScale scale_ab = tCrSFA_local(_0{}); + scale_if_needed(accum_local, accumulation, scale_ab); + } + if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) { + scale_if_needed(accum_local, accumulation, tCrSFA_local); + } + if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) { + scale_if_needed(accum_local, accumulation, tCrSFB_local); + } + if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile > 1) { + scale_if_needed(accum_local, accumulation, tCrSFA_local, tCrSFB_local); + } + if constexpr (ScalePromotionInterval != 4) { + // residues only exists when granularity is not the finnest + if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) { + ElementBlockScale scale_ab = tCrSFA_local(_0{}); + accumulation.scale_residue_if_needed(accum_local, scale_ab); + } + if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) { + accumulation.scale_residue_if_needed(accum_local, tCrSFA_local); + } + if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) { + accumulation.scale_residue_if_needed(accum_local, tCrSFB_local); + } + if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) { + accumulation.scale_residue_if_needed(accum_local, tCrSFA_local, tCrSFB_local); + } + } + + warpgroup_fence_operand(accumulation()); + } // end for (m_split) + } } /// Perform a Consumer Epilogue to release all buffers CUTLASS_DEVICE void mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) { - // The pipeline is not released in the first iteration - smem_pipe_release.advance(k_tile_count - 1); - pipeline.consumer_release(smem_pipe_release); } }; diff --git a/include/cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized_fp8.hpp b/include/cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized_fp8.hpp index d993d9a1..6f58a20f 100644 --- a/include/cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized_fp8.hpp +++ b/include/cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized_fp8.hpp @@ -586,8 +586,8 @@ struct CollectiveMma< int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count); tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; - - GmmaFP8Accumulation accumulation(accum, mainloop_params.mma_promotion_interval, size<2>(tCrA)); + auto accm_temp = cute::make_fragment_like(accum); + GmmaFP8Accumulation accumulation(accm_temp, mainloop_params.mma_promotion_interval, size<2>(tCrA)); warpgroup_fence_operand(accumulation()); CUTLASS_PRAGMA_UNROLL for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0; --k_tile_prologue) @@ -614,7 +614,7 @@ struct CollectiveMma< warpgroup_commit_batch(); - accumulation.promote_if_needed(); + accumulation.promote_if_needed(accum); ++smem_pipe_read; } @@ -652,7 +652,7 @@ struct CollectiveMma< warpgroup_wait(); warpgroup_fence_operand(accumulation()); - accumulation.promote_if_needed(); + accumulation.promote_if_needed(accum); // UNLOCK smem_pipe_release, done _computing_ on it pipeline.consumer_release(smem_pipe_release); @@ -662,7 +662,7 @@ struct CollectiveMma< ++smem_pipe_release; } - accumulation.promote_residue_if_needed(); + accumulation.promote_residue_if_needed(accum); warpgroup_fence_operand(accumulation()); } diff --git a/include/cutlass/gemm/group_array_problem_shape.hpp b/include/cutlass/gemm/group_array_problem_shape.hpp index 400f7e6b..6d198e17 100644 --- a/include/cutlass/gemm/group_array_problem_shape.hpp +++ b/include/cutlass/gemm/group_array_problem_shape.hpp @@ -44,6 +44,7 @@ #include #endif + //////////////////////////////////////////////////////////////////////////////////////////////////// namespace cutlass::gemm { @@ -57,6 +58,7 @@ struct GroupProblemShape { UnderlyingProblemShape* problem_shapes = nullptr; UnderlyingProblemShape const* host_problem_shapes = nullptr; + CUTLASS_HOST_DEVICE int32_t groups() const { return num_groups; } @@ -79,16 +81,56 @@ struct GroupProblemShape { } }; -template +template struct MoEProblemShape { + using UnderlyingProblemShape = ProblemShape_; - using MaxProblemShape = MaxProblemShape_; + static_assert(rank(UnderlyingProblemShape{}) == 3, "ProblemShape{} should be "); + + int32_t max_m = 0; + int32_t max_n = 0; + int32_t max_k = 0; + int32_t num_groups = 0; + int32_t* tokens_per_expert = nullptr; + int32_t* tokens_per_expert_host = nullptr; + + CUTLASS_HOST_DEVICE + int32_t groups() const { return num_groups; } + + CUTLASS_HOST_DEVICE + UnderlyingProblemShape const + get_problem_shape(int32_t group_idx=0) const { + + UnderlyingProblemShape expert_problem_dims; + assert(tokens_per_expert != nullptr); //tokens_per_expert should not be null + if (group_idx < num_groups) { // add check on the can_implement + expert_problem_dims = {max_m, tokens_per_expert[group_idx], max_k}; + } + + return expert_problem_dims; + } + + // Function returns max problem shape if tokens_per_expert host is unavailable. + // Returns host problem shape if tokens_per_expert host is available. + CUTLASS_HOST_DEVICE + UnderlyingProblemShape const + get_host_problem_shape(int32_t group_idx=0) const { + UnderlyingProblemShape expert_problem_dims = {max_m, max_n, max_k}; + assert(tokens_per_expert_host != nullptr); //tokens_per_expert_host should not be null + if (group_idx < num_groups) { + expert_problem_dims = {max_m, tokens_per_expert_host[group_idx], max_k}; + } + return expert_problem_dims; + } + + CUTLASS_HOST_DEVICE + bool + is_host_problem_shape_available() const { + return tokens_per_expert_host != nullptr; + } - UnderlyingProblemShape problem_shape; - MaxProblemShape max_problem_shape; }; - template class ArrayProblemShape { public: @@ -135,8 +177,9 @@ namespace detail { template struct is_moe_problem_shape : cute::false_type {}; -template -struct is_moe_problem_shape> : cute::true_type {}; + +template +struct is_moe_problem_shape> : cute::true_type {}; } 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 c03d67c0..ad510369 100644 --- a/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp +++ b/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp @@ -842,13 +842,12 @@ public: ); auto work_tile_info = [&] () { - if constexpr (!IsSchedDynamicPersistent) { - // 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(); - } + // 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 9d5e5b4e..6f853f81 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 @@ -649,18 +649,21 @@ public: transform2mma_pipeline.init_masks(cluster_shape); mma2accum_pipeline.init_masks(cluster_shape); + // 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); - typename TileScheduler::WorkTileInfo work_tile_info = scheduler.initial_work_tile_info(cluster_shape); + // Ensure memory ops in this kernel are not done prior to completion of dependent grids. + cutlass::arch::wait_on_dependent_grids(); + + 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); // Optionally append 1s until problem shape is rank-4 in case it is only rank-3 (MNK) auto problem_shape_MNKL = append<4>(problem_shape.get_problem_shape(work_tile_info.L_idx), 1); - // Allocate accumulators - auto acc_shape = collective_mainloop.partition_accumulator_shape(); - // NOTE: we can assume the tmem buf starts at zero since we allocate all tmem in this kernel auto bulk_tmem = TiledMma::make_fragment_C(append(acc_shape, Int{})); 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 ee670016..5c8cb84e 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 @@ -726,11 +726,6 @@ public: mainloop_ab_pipeline.init_masks(cluster_shape, block_id_in_cluster); accumulator_pipeline.init_masks(cluster_shape, block_id_in_cluster); - // 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); - // // TMEM "Allocation" // @@ -740,6 +735,15 @@ 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(); + + 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); + pipeline_init_wait(cluster_size); if constexpr (IsGroupedGemmKernel) { diff --git a/include/cutlass/gemm/kernel/sm100_gemm_mixed_tma_cpasync_warpspecialized.hpp b/include/cutlass/gemm/kernel/sm100_gemm_mixed_tma_cpasync_warpspecialized.hpp index 99da60bf..701149c2 100644 --- a/include/cutlass/gemm/kernel/sm100_gemm_mixed_tma_cpasync_warpspecialized.hpp +++ b/include/cutlass/gemm/kernel/sm100_gemm_mixed_tma_cpasync_warpspecialized.hpp @@ -83,20 +83,20 @@ public: CUTLASS_HOST_DEVICE static auto get_problem_shape_gemm(ProblemShape const& shape) { if constexpr (IsGroupedGemmKernel) { - return shape.max_problem_shape; + auto problem_shape_MNK = shape.get_host_problem_shape(0); //gets the maximum problem shape here + auto problem_shape_MNKL = append<4>(problem_shape_MNK, shape.groups()); //appends num_groups to it + return problem_shape_MNKL; } else { return shape; } } + CUTLASS_HOST_DEVICE static auto get_problem_shape_scheduler(ProblemShape const& shape) { - if constexpr (IsMoEScheduler) { + if constexpr (IsMoEScheduler) { return shape; } - else if constexpr (IsGroupedGemmKernel) { - return shape.problem_shape; - } else { return shape; } @@ -106,7 +106,7 @@ public: CUTLASS_HOST_DEVICE static auto get_effective_shape(ProblemShape const& shape, WorkTileInfo const& work_tile_info) { if constexpr (IsGroupedGemmKernel) { - return append<4>(shape.problem_shape.get_problem_shape(work_tile_info.L_idx), Int<1>{}); + return append<4>(shape.get_problem_shape(work_tile_info.L_idx), Int<1>{}); } else { return append<4>(shape, Int<1>{}); @@ -114,10 +114,9 @@ public: } using ProblemShapeGemm = decltype(get_problem_shape_gemm(ProblemShape{})); - using ProblemShapeScheduler = decltype(get_problem_shape_scheduler(ProblemShape{})); static_assert(rank(ProblemShapeGemm{}) == 3 or rank(ProblemShapeGemm{}) == 4, - "ProblemShapeGemm{} should be or "); + "ProblemShape{} should be or "); static constexpr bool IsGdcEnabled = false; // Mainloop derived types using CollectiveMainloop = CollectiveMainloop_; @@ -160,7 +159,7 @@ public: static_assert(size(AtomThrShapeMNK{}) == 1, "Lower alignment kernel only supports 1x1x1 cluster shape."); using TileSchedulerTag = cute::conditional_t; using TileScheduler = typename detail::TileSchedulerSelector< - TileSchedulerTag, ArchTag, CtaShape_MNK, ClusterShape, SchedulerPipelineStageCount, ProblemShapeScheduler>::Scheduler; + TileSchedulerTag, ArchTag, CtaShape_MNK, ClusterShape, SchedulerPipelineStageCount, ProblemShape>::Scheduler; using TileSchedulerArguments = typename TileScheduler::Arguments; using TileSchedulerParams = typename TileScheduler::Params; @@ -259,7 +258,6 @@ public: GemmUniversalMode mode{}; ProblemShape problem_shape{}; ProblemShapeGemm problem_shape_gemm{}; - ProblemShapeScheduler problem_shape_scheduler{}; MainloopParams mainloop{}; EpilogueParams epilogue{}; KernelHardwareInfo hw_info{}; @@ -289,18 +287,20 @@ public: Params to_underlying_arguments(Arguments const& args, void* workspace) { (void) workspace; - // auto problem_shape = args.problem_shape; - // auto problem_shape_MNKL = append<4>(problem_shape, 1); - + auto problem_shape = args.problem_shape; auto problem_shape_gemm = get_problem_shape_gemm(args.problem_shape); - auto problem_shape_scheduler = get_problem_shape_scheduler(args.problem_shape); + // Get SM count if needed, otherwise use user supplied SM count int sm_count = args.hw_info.sm_count; - if (sm_count != 0) { + if (IsGroupedGemmKernel && sm_count <= 0) { + CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n" + " For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count."); + sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id); + } + else if (!IsGroupedGemmKernel && sm_count != 0) { CUTLASS_TRACE_HOST(" WARNING: SM100 tile scheduler does not allow for user specified SM counts.\n" " To restrict a kernel's resource usage, consider using CUDA driver APIs instead (green contexts)."); - sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id); } CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count); @@ -313,25 +313,24 @@ public: // Epilogue void* epilogue_workspace = workspace_ptr + workspace_offset; - workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue); + workspace_offset += CollectiveEpilogue::get_workspace_size(problem_shape, args.epilogue); workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); void* mainloop_workspace = nullptr; // Tile scheduler void* scheduler_workspace = workspace_ptr + workspace_offset; - workspace_offset += TileScheduler::template get_workspace_size( - args.scheduler, problem_shape_scheduler, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs); + workspace_offset += TileScheduler::template get_workspace_size( + args.scheduler, problem_shape, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs); workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); TileSchedulerParams scheduler; if constexpr (IsGroupedGemmKernel) { scheduler = TileScheduler::to_underlying_arguments( - problem_shape_scheduler, TileShape{}, AtomThrShapeMNK{}, ClusterShape{}, + problem_shape, TileShape{}, AtomThrShapeMNK{}, ClusterShape{}, args.hw_info, args.scheduler, scheduler_workspace); } else { - auto problem_shape = args.problem_shape; auto problem_shape_MNKL = append<4>(problem_shape, 1); scheduler = TileScheduler::to_underlying_arguments( @@ -344,7 +343,6 @@ public: args.mode, args.problem_shape, problem_shape_gemm, - problem_shape_scheduler, CollectiveMainloop::to_underlying_arguments(problem_shape_gemm, args.mainloop, mainloop_workspace), CollectiveEpilogue::to_underlying_arguments(problem_shape_gemm, args.epilogue, epilogue_workspace), hw_info, @@ -358,8 +356,7 @@ public: if constexpr (IsGroupedGemmKernel) { implementable &= args.mode == GemmUniversalMode::kGrouped; - implementable &= rank(ProblemShapeGemm{}) == 4; - implementable &= rank(typename ProblemShape::UnderlyingProblemShape::UnderlyingProblemShape{}) == 3; + implementable &= rank(typename ProblemShape::UnderlyingProblemShape{}) == 3; } else { implementable &= (args.mode == GemmUniversalMode::kGemm) or @@ -387,15 +384,14 @@ public: size_t workspace_size = 0; auto problem_shape_gemm = get_problem_shape_gemm(args.problem_shape); - auto problem_shape_scheduler = get_problem_shape_scheduler(args.problem_shape); // Epilogue workspace_size += CollectiveEpilogue::get_workspace_size(problem_shape_gemm, args.epilogue); workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment); // Tile scheduler - workspace_size += TileScheduler::template get_workspace_size( - args.scheduler, problem_shape_scheduler, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs); + workspace_size += TileScheduler::template get_workspace_size( + args.scheduler, args.problem_shape, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs); workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment); return workspace_size; @@ -409,7 +405,6 @@ public: size_t workspace_offset = 0; auto problem_shape_gemm = get_problem_shape_gemm(args.problem_shape); - auto problem_shape_scheduler = get_problem_shape_scheduler(args.problem_shape); // Epilogue status = CollectiveEpilogue::initialize_workspace(problem_shape_gemm, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter); @@ -421,10 +416,10 @@ public: } // Tile scheduler - status = TileScheduler::template initialize_workspace( - args.scheduler, workspace_ptr + workspace_offset, stream, problem_shape_scheduler, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs, cuda_adapter); - workspace_offset += TileScheduler::template get_workspace_size( - args.scheduler, problem_shape_scheduler, args.hw_info, NumFixupBarriers); + status = TileScheduler::template initialize_workspace( + args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs, cuda_adapter); + workspace_offset += TileScheduler::template get_workspace_size( + args.scheduler, args.problem_shape, args.hw_info, NumFixupBarriers); workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); if (status != Status::kSuccess) { return status; @@ -441,14 +436,14 @@ public: if constexpr (IsGroupedGemmKernel) { grid_shape = TileScheduler::get_grid_shape( params.scheduler, - params.problem_shape_scheduler, + params.problem_shape, TileShape{}, AtomThrShapeMNK{}, cluster_shape, params.hw_info); } else { - auto problem_shape_MNKL = append<4>(params.problem_shape_scheduler, 1); + auto problem_shape_MNKL = append<4>(params.problem_shape, 1); grid_shape = TileScheduler::get_grid_shape( params.scheduler, problem_shape_MNKL, @@ -505,8 +500,6 @@ public: // Do we load source tensor C or other aux inputs bool is_epi_load_needed = collective_epilogue.is_producer_load_needed(); - // printf("is_epi_load_needed = %d", (int)is_epi_load_needed); - IsParticipant is_participant = { (warp_category == WarpCategory::MMA) && is_mma_leader_cta, // mma (warp_category == WarpCategory::Sched) && is_first_cta_in_cluster, // sched @@ -654,9 +647,6 @@ public: // // TMEM "Allocation" // - // auto acc_shape = collective_mainloop.partition_accumulator_shape(); - // auto bulk_tmem = TiledMma::make_fragment_C(append(acc_shape, - // Int{})); auto tmem_storage = collective_mainloop.template init_tmem_tensors(EpilogueTile{}); // @@ -666,10 +656,10 @@ public: // Synchronization call. Blocks until barriers are initialized in shared memory. pipeline_init_wait(cluster_size); - // __syncwarp(); - // if (threadIdx.x % 32 == 0) { - // printf("warp %d start\n", warp_idx); - // } + if (not work_tile_info.is_valid()) { + // When problem shapes are only on device, the grid launched may be larger than the total number of blocks across groups + return; + } if (is_participant.main_load_tma) { // Ensure that the prefetched kernel does not touch @@ -689,7 +679,6 @@ public: // Get the number of K tiles to compute for this work as well as the starting K tile offset of the work. auto k_tile_iter = scheduler.get_k_tile_iterator(work_tile_info, effective_shape, CtaShape_MNK{}, k_tiles); auto k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, effective_shape, CtaShape_MNK{}); - // auto k_tile_prologue = min(MainloopPipeline::Stages, k_tile_count); auto [mainloop_producer_state_next_, unused_] = collective_mainloop.load_tma( 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 6708a88c..0411b585 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 @@ -734,11 +734,6 @@ public: mainloop_ab_pipeline.init_masks(cluster_shape); mainloop_sf_pipeline.init_masks(cluster_shape); accumulator_pipeline.init_masks(cluster_shape); - // 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); - // // TMEM "Allocation" // @@ -749,6 +744,15 @@ 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(); + + 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); + pipeline_init_wait(cluster_size); if constexpr (IsGroupedGemmKernel) { 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 ec5cd4d0..c69defbf 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 @@ -48,6 +48,7 @@ #include "cutlass/trace.h" #include "cutlass/gemm/kernel/sm90_tile_scheduler.hpp" #include "cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp" +#include "cutlass/arch/grid_dependency_control.h" /////////////////////////////////////////////////////////////////////////////// @@ -418,7 +419,7 @@ public: // Any Tensor Op MMA Atom in the ISA is arch conditional. #if ! defined(ENABLE_SM90_KERNEL_LEVEL) - printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else // Preconditions @@ -581,6 +582,9 @@ public: // Wait for all thread blocks in the Cluster cluster_wait_fn(); + // Ensure memory ops in this kernel are not done prior to completion of dependent grids. + cutlass::arch::wait_on_dependent_grids(); + 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 fd7ff603..d380b2cd 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 @@ -48,6 +48,7 @@ #include "cutlass/trace.h" #include "cutlass/gemm/kernel/sm90_tile_scheduler.hpp" #include "cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp" +#include "cutlass/arch/grid_dependency_control.h" /////////////////////////////////////////////////////////////////////////////// @@ -430,7 +431,7 @@ public: // Any Tensor Op MMA Atom in the ISA is arch conditional. #if ! defined(ENABLE_SM90_KERNEL_LEVEL) - printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else // Preconditions @@ -596,6 +597,9 @@ public: // Wait for all thread blocks in the Cluster cluster_wait_fn(); + // Ensure memory ops in this kernel are not done prior to completion of dependent grids. + cutlass::arch::wait_on_dependent_grids(); + 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_tma.hpp b/include/cutlass/gemm/kernel/sm90_gemm_tma.hpp index 2292d7e4..d023c1d5 100644 --- a/include/cutlass/gemm/kernel/sm90_gemm_tma.hpp +++ b/include/cutlass/gemm/kernel/sm90_gemm_tma.hpp @@ -42,6 +42,8 @@ #include "cutlass/gemm/kernel/sm90_tile_scheduler.hpp" #include "cutlass/gemm/kernel/tile_scheduler.hpp" #include "cutlass/trace.h" +#include "cutlass/arch/grid_dependency_control.h" + #include "cute/tensor.hpp" /////////////////////////////////////////////////////////////////////////////// @@ -204,7 +206,7 @@ public: // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a. #if ! defined(__CUDA_ARCH_FEAT_SM90_ALL) - printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); #else // Preconditions @@ -261,6 +263,9 @@ public: auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA)); auto k_tile_count = size<2>(gA); + // Ensure memory ops in this kernel are not done prior to completion of dependent grids. + cutlass::arch::wait_on_dependent_grids(); + // Perform the collective scoped MMA CollectiveMainloop collective_mma; collective_mma( diff --git a/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp b/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp index 5b558005..8370ccb9 100644 --- a/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp +++ b/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp @@ -277,7 +277,7 @@ public: // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a. #if ! defined(ENABLE_SM90_KERNEL_LEVEL) - printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); #else enum class WarpGroupRole { diff --git a/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp b/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp index d398d1f2..5c04259b 100644 --- a/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp +++ b/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp @@ -132,7 +132,21 @@ public: static constexpr int RegsPerThread = size<0>(TileShape{}) * size<1>(TileShape{}) / NumMMAThreads * sizeof(ElementAccumulator) / sizeof(uint32_t); - static constexpr bool HeavyRegisterPressure = RegsPerThread >= 208; + + // Detect if this is SM120 blockscaled kernel which hits high register pressure + // on smaller tiles (e.g. 256x128 registers per thread) + template + struct is_blockscaled : cute::false_type {}; + + template + struct is_blockscaled> + : cute::true_type {}; + + static constexpr bool IsBlockScaled = is_blockscaled::value; + + static constexpr bool HeavyRegisterPressure = + IsBlockScaled ? (RegsPerThread >= 128) : (RegsPerThread >= 208); + static constexpr uint32_t LoadRegisterRequirement = !HeavyRegisterPressure ? 40 : 24; static constexpr uint32_t MmaRegisterRequirement = !HeavyRegisterPressure ? 232 : 240; @@ -349,7 +363,7 @@ public: // Any Tensor Op MMA Atom in the ISA is arch conditional. #if ! defined(ENABLE_SM90_KERNEL_LEVEL) - printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else // Preconditions diff --git a/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp b/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp index 1326f390..073f3a50 100644 --- a/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp +++ b/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp @@ -361,7 +361,7 @@ public: // Any Tensor Op MMA Atom in the ISA is arch conditional. #if ! defined(ENABLE_SM90_KERNEL_LEVEL) - printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else // Preconditions diff --git a/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized.hpp b/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized.hpp index e7cafde5..71411d56 100644 --- a/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized.hpp +++ b/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized.hpp @@ -42,6 +42,8 @@ #include "cutlass/gemm/kernel/sm90_tile_scheduler.hpp" #include "cutlass/pipeline/pipeline.hpp" #include "cute/tensor.hpp" +#include "cutlass/arch/grid_dependency_control.h" + /////////////////////////////////////////////////////////////////////////////// namespace cutlass::gemm::kernel { @@ -227,7 +229,7 @@ public: // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a. #if ! defined(__CUDA_ARCH_FEAT_SM90_ALL) - printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); #else enum class WarpGroupRole { @@ -343,6 +345,9 @@ public: auto k_residue = K - size<1>(gA) * size<2>(gA); // K - BLK_K * k_coord_max auto residue_mnk = make_tuple(m_max_coord, n_max_coord, k_residue); + // Ensure memory ops in this kernel are not done prior to completion of dependent grids. + cutlass::arch::wait_on_dependent_grids(); + collective_mainloop.load( mainloop_pipeline, mainloop_pipe_producer_state, diff --git a/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_cooperative.hpp b/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_cooperative.hpp index 1d35ff2d..11ccbfc2 100644 --- a/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_cooperative.hpp +++ b/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_cooperative.hpp @@ -42,6 +42,7 @@ #include "cutlass/gemm/kernel/tile_scheduler.hpp" #include "cutlass/pipeline/pipeline.hpp" #include "cute/tensor.hpp" +#include "cutlass/arch/grid_dependency_control.h" /////////////////////////////////////////////////////////////////////////////// @@ -265,7 +266,7 @@ public: // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a. #if ! defined(__CUDA_ARCH_FEAT_SM90_ALL) - printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); #else static_assert(cute::rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>."); @@ -386,6 +387,9 @@ public: auto k_residue = K - size<1>(gA) * size<2>(gA); // K - BLK_K * k_coord_max auto residue_mnk = make_tuple(m_max_coord, n_max_coord, k_residue); + // Ensure memory ops in this kernel are not done prior to completion of dependent grids. + cutlass::arch::wait_on_dependent_grids(); + collective_mainloop.load( mainloop_pipeline, mainloop_pipe_producer_state, diff --git a/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_pingpong.hpp b/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_pingpong.hpp index be086f0c..57fcd32e 100644 --- a/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_pingpong.hpp +++ b/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_pingpong.hpp @@ -45,6 +45,8 @@ #include "cutlass/trace.h" #include "cute/tensor.hpp" +#include "cutlass/arch/grid_dependency_control.h" + /////////////////////////////////////////////////////////////////////////////// namespace cutlass::gemm::kernel { @@ -271,7 +273,7 @@ public: // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a. #if ! defined(__CUDA_ARCH_FEAT_SM90_ALL) - printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); + CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); #else // Preconditions @@ -409,6 +411,10 @@ public: auto k_residue = K - size<1>(gA) * size<2>(gA); // K - BLK_K * k_coord_max auto residue_mnk = make_tuple(m_max_coord, n_max_coord, k_residue); + // Ensure memory ops in this kernel are not done prior to completion of dependent grids. + cutlass::arch::wait_on_dependent_grids(); + + collective_mainloop.load( mainloop_pipeline, mainloop_pipe_producer_state, diff --git a/include/cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp b/include/cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp index 577c68c3..4332c6ad 100644 --- a/include/cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp +++ b/include/cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp @@ -378,13 +378,40 @@ private: } // Construct a sign bit mask for handling negative zeros - ElementAMmaRawUnit sign_mask = ElementAMmaRawUnit{ 0 }; - if constexpr (has_negative_zero_v) { - ElementAMmaRawUnit one_sign_mask = static_cast(~(ElementAMmaRawUnit{ 1 } << (cute::sizeof_bits_v - 1))); - for (int i = 0; i < sizeof(ElementAMmaRawUnit) / sizeof(ElementAUint); ++i) { - sign_mask = static_cast((int32_t)sign_mask | (int32_t)one_sign_mask << (i * cute::sizeof_bits_v)); + // Compute the mask value at compile time, then construct ElementAMmaRawUnit from it + // Case 1: float_e2m1_t (4-bit), uint4_t container, ElemsARawPerElementAMmaRaw 1 element + // - Result: negzero_mask_value = 0b0111 (stored in uint8_t as 0b0000_0111) + // + // Case 2: float_e2m1_t (4-bit), uint8_t container, ElemsARawPerElementAMmaRaw 2 elements + // - Result: negzero_mask_value = 0b0111_0111 + // + // Case 3: float_e4m3_t (8-bit), uint8_t container, ElemsARawPerElementAMmaRaw 1 element + // - Result: negzero_mask_value = 0b0111_1111 = 0x7F + // Note: Lambda returns uint32_t (constexpr-compatible) instead of ElementAMmaRawUnit (non-literal type) + constexpr uint32_t negzero_mask_value = []() constexpr -> uint32_t { + constexpr int ElementAMmaRawNumBits = cute::sizeof_bits_v; + constexpr int ElementANumBits = cute::sizeof_bits_v; + + if constexpr (has_negative_zero_v) { + // Create mask for one ElementA: all bits set except the sign bit (MSB) + // ElementANumBits = 4: (1 << 3) - 1 = 0b0111 + // ElementANumBits = 8: (1 << 7) - 1 = 0b0111_1111 + constexpr uint32_t ElementASignMask = (1u << (ElementANumBits - 1)) - 1; + + // Replicate the single-element mask across all packed elements + if constexpr (ElemsARawPerElementAMmaRaw == 1) { + return ElementASignMask; + } + else if constexpr (ElemsARawPerElementAMmaRaw == 2) { + return (ElementASignMask << ElementANumBits) | ElementASignMask; + } } - } + // No negative zero: return all bits set to 1 (no masking needed) + return (1u << ElementAMmaRawNumBits) - 1; + }(); + + // Construct ElementAMmaRawUnit from the compile-time computed mask value + const ElementAMmaRawUnit negzero_mask_out_sign_mask = ElementAMmaRawUnit{negzero_mask_value}; // * Compress // cACsAC is always row major order @@ -410,17 +437,21 @@ private: // * Find None Zero Element Idx within Chunk CUTE_UNROLL for (int elt_log_idx = 0; elt_log_idx < OneChunkSizeA{}; ++elt_log_idx) { - ElementAMmaRawUnit elem_A = tAsA[elt_log_idx]; + // Iterate through all ElementAMma within one logical chunk + ElementAMmaRawUnit tAsA_i = tAsA[elt_log_idx]; - // Handle negative 0 - ElementAMmaRawUnit masked_elem_A = elem_A; + // Mask off the signed bit s.t. negative zero is same as positive zero + ElementAMmaRawUnit tAsA_i_negzero_masked_out = tAsA_i; if constexpr (has_negative_zero_v) { - masked_elem_A = elem_A & sign_mask; + // For sub-bytes, LSB will contain valid bits. + // e.g. for float_e2m1_t, tAsA_i is stored in uint8_t with 0x0000yyyy where yyyy denote valid bits. + tAsA_i_negzero_masked_out = tAsA_i & negzero_mask_out_sign_mask; } - if (masked_elem_A != ElementAMmaRawUnit{0}) { + // Record this ElmentAMma if it's none zero + if (tAsA_i_negzero_masked_out != ElementAMmaRawUnit{0}) { non_zero_elt_log_idx[non_zero_cnt] = elt_log_idx; - tACsAC[non_zero_cnt] = elem_A; + tACsAC[non_zero_cnt] = tAsA_i; non_zero_cnt++; } } diff --git a/media/docs/cpp/ide_setup.md b/media/docs/cpp/ide_setup.md index bad80bba..95933d5c 100644 --- a/media/docs/cpp/ide_setup.md +++ b/media/docs/cpp/ide_setup.md @@ -23,7 +23,7 @@ and you might see faster responses and more stable performance with clangd. 1. Enter "C/C++" to filter results 1. Select "C/C++ Edit Configurations (UI)" (or "... (JSON)" if you feel like editing the raw JSON) 1. View the documentation for these settings - [here](https://code.visualstudio.com/docs/cpp/c-cpp-properties-schema-reference) + [here](https://code.visualstudio.com/docs/cpp/customize-cpp-settings) 1. Edit "Include Path" to set up **include paths**. For CUTLASS, this includes the following: * `${workspaceFolder}/include` * `${workspaceFolder}/tools/util/include` diff --git a/media/docs/pythonDSL/cute_dsl.rst b/media/docs/pythonDSL/cute_dsl.rst index a7f53587..1c129a77 100644 --- a/media/docs/pythonDSL/cute_dsl.rst +++ b/media/docs/pythonDSL/cute_dsl.rst @@ -17,3 +17,4 @@ CuTe DSL Debugging with the DSL Autotuning with the DSL Educational Notebooks + Compile with TVM FFI diff --git a/media/docs/pythonDSL/cute_dsl_api/changelog.rst b/media/docs/pythonDSL/cute_dsl_api/changelog.rst index 4de8d6e0..49a66caa 100644 --- a/media/docs/pythonDSL/cute_dsl_api/changelog.rst +++ b/media/docs/pythonDSL/cute_dsl_api/changelog.rst @@ -2,29 +2,35 @@ Changelog for CuTe DSL API changes ====================================== -`4.3.0 `_ (2025-10-07) +`4.3.0 `_ (2025-10-20) ============================================================================== * Debuggability improvements: + - Supported source location tracking for DSL APIs - - Supported dumping PTX and SASS code -* Remove deprecated ``cutlass._utils.SMEM_CAPACITY[""]`` and ``cutlass.utils.ampere_helpers`` -* Support calling nested functions without capturing variables inside dynamic control flow -* Replace usage of ``cute.arch.barrier`` in examples with corresponding APIs in ``pipeline`` + - Supported dumping PTX and CUBIN + +* Removed deprecated ``cutlass._utils.SMEM_CAPACITY[""]`` and ``cutlass.utils.ampere_helpers`` +* Supported calling nested functions without capturing variables inside dynamic control flow +* Replaced usage of ``cute.arch.barrier`` in examples with corresponding APIs in ``pipeline`` + - Use ``pipeline.sync`` for simple cases like synchronizing the whole CTA - Use ``pipeline.NamedBarrier`` to customize barriers with different participating threads and barrier id + * Added new APIs ``repeat`` and ``repeat_as_tuple`` -* Added new APIs ``make_rmem_tensor`` to replace ``make_fragment`` with better naming +* Added new APIs ``make_rmem_tensor`` to create tensor in register memory (replace ``make_fragment`` with better naming) * Added new APIs ``make_rmem_tensor_like`` which create rmem tensor from a tensor using the same shape with compact col-major strides * Added ``TmemAllocator`` for allocating tensor memory * Updated ``SmemAllocator.allocate`` to support allocation of a single scalar value * Fixed ``TensorSSA.reduce`` to support static value as initial value * Updated docstring for following APIs to be more concise and easier to understand: + - ``make_layout_tv`` - ``is_static`` - ``PipelineAsync`` - ``SmemAllocator`` -* Fixed documentation for ``pipeline``, ``utils`` and ``cute.math`` + +* Fixed documentation for ``pipeline``, ``utils`` and ``cute.math`` (``cute.math`` is part of top level documentation) `4.2.0 `_ (2025-09-10) @@ -32,7 +38,9 @@ Changelog for CuTe DSL API changes * Added back ``cute.make_tiled_copy`` per the request from community * Added support for explicit and implicit broadcast in ``TensorSSA`` + - ``cutlass.cute.TensorSSA``: support ``broadcast_to`` and implicit broadcasting for binary operations. + * Supported printing ``TensorSSA`` value in ``cutlass.cute.print_tensor`` * Updated ``cute.gemm`` to support all dispatch patterns and improved checks for illegal inputs * Introduced automatic kernel smem usage calculation for launch config. diff --git a/media/docs/pythonDSL/cute_dsl_api/cute.rst b/media/docs/pythonDSL/cute_dsl_api/cute.rst index f64e6b03..92c282c8 100644 --- a/media/docs/pythonDSL/cute_dsl_api/cute.rst +++ b/media/docs/pythonDSL/cute_dsl_api/cute.rst @@ -15,3 +15,4 @@ cutlass.cute :hidden: cute_arch + cute_runtime diff --git a/media/docs/pythonDSL/cute_dsl_api/cute_runtime.rst b/media/docs/pythonDSL/cute_dsl_api/cute_runtime.rst new file mode 100644 index 00000000..bf64e11b --- /dev/null +++ b/media/docs/pythonDSL/cute_dsl_api/cute_runtime.rst @@ -0,0 +1,16 @@ +.. _cute_runtime: + +Runtime +======= + +Description + +API documentation +----------------- + +.. automodule:: cutlass.cute.runtime + :members: + :undoc-members: + :show-inheritance: + :special-members: __init__ + :private-members: diff --git a/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst b/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst new file mode 100644 index 00000000..ad70f2a2 --- /dev/null +++ b/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst @@ -0,0 +1,381 @@ +.. _compile_with_tvm_ffi: +.. |DSL| replace:: CuTe DSL + +Compile with TVM FFI +==================== + +Apache TVM FFI is an open ABI and FFI for machine learning systems. More information can be found in the `official documentation `_. + +To install TVM FFI, you can run the following command: + +.. code-block:: bash + + pip install apache-tvm-ffi + # optional package for improved torch tensor calling performance + pip install torch-c-dlpack-ext + +In |DSL|, TVM FFI can be enabled as an option for JIT-compiled functions. Using TVM FFI can lead to faster JIT function invocation and provides better interoperability with machine learning frameworks (e.g., directly take ``torch.Tensor`` as arguments). + + +Enable Apache TVM FFI in |DSL| +------------------------------ + +First, install the ``tvm-ffi`` package by following its `installation guide `_. + +There are two ways to enable TVM FFI in |DSL|: + +1. Use the ``options`` argument in ``cute.compile`` to specify the TVM FFI option. For example: + +.. code-block:: python + + # Assuming you have defined a function `add` decorated with @cute.jit + def example_compile(): + a_torch = torch.randn(10, 20, 30).to(torch.float16) + b_torch = torch.randn(10, 20, 30).to(torch.float16) + a_cute = cute.runtime.from_dlpack(a_torch, enable_tvm_ffi=True).mark_layout_dynamic() + b_cute = cute.runtime.from_dlpack(b_torch, enable_tvm_ffi=True).mark_layout_dynamic() + + compiled_add = cute.compile(add, a_cute, b_cute, options="--enable-tvm-ffi") + + +Note that the object returned by ``cute.compile`` is a Python function specific to TVM FFI. + +2. Alternatively, you can enable TVM FFI globally by setting the environment variable ``CUTE_DSL_ENABLE_TVM_FFI=1``. Please note that this setting will apply to all JIT compilations within the environment. + +Fake tensor for compilation +--------------------------- + +The TVM FFI function accepts DLPack-compatible tensors as arguments, such as those from torch or jax. +However, during compilation, it is necessary to specify the tensors' dynamic properties in |DSL|. +To clearly distinguish between the compilation phase and runtime, +|DSL| provides a "fake tensor" that can be used for compilation. For example: + +.. code-block:: python + + import cutlass.cute as cute + import torch + + @cute.kernel + def device_add_one(a: cute.Tensor, b: cute.Tensor): + threads_per_block = 128 + cta_x_, _, _ = cute.arch.block_idx() + tid_x, _, _ = cute.arch.thread_idx() + tid = cta_x_ * threads_per_block + tid_x + if tid < a.shape[0]: + b[tid] = a[tid] + 1.0 + + @cute.jit + def add_one(a: cute.Tensor, b: cute.Tensor): + n = a.shape[0] + threads_per_block = 128 + blocks = (n + threads_per_block - 1) // threads_per_block + device_add_one(a, b).launch( + grid=(blocks, 1, 1), + block=(threads_per_block, 1, 1), + ) + + def example_add_one(): + n = cute.sym_int() + a_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,)) + b_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,)) + # compile the kernel with "--enable-tvm-ffi" option and example input tensors + compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi") + # now compiled_add_one is a TVM-FFI function that can be called with torch.Tensor as input + a_torch = torch.arange(10, dtype=torch.float32, device="cuda") + b_torch = torch.empty(10, dtype=torch.float32, device="cuda") + compiled_add_one(a_torch, b_torch) + print("result of b_torch after compiled_add_one(a_torch, b_torch)") + print(b_torch) + + +The fake tensor is a placeholder that mimics the interface of a real tensor but does not hold real data or allow indexing. +It is used in compilation or testing scenarios where only shape/type/layout information is needed. +All attempts to access or mutate data will raise errors. + +``cute.Tensor`` adapter for TVM FFI +----------------------------------- + +To adapt the ``cute.Tensor`` to the TVM FFI function, you can use the ``cute.runtime.from_dlpack`` function with the ``enable_tvm_ffi=True`` option or the environment variable ``CUTE_DSL_ENABLE_TVM_FFI=1``. For example: + +.. code-block:: python + + def example_from_dlpack(): + a_cute = cute.runtime.from_dlpack(a_torch, enable_tvm_ffi=True).mark_layout_dynamic() + b_cute = cute.runtime.from_dlpack(b_torch, enable_tvm_ffi=True).mark_layout_dynamic() + + compiled_add_one(a_cute, b_cute) + +Note that because the ``cute.runtime.from_dlpack`` function performs an explicit DLPack conversion, it is less efficient than passing the ``torch.Tensor`` directly. +You can also use ``cute.Tensor`` as an argument hint for ``cute.compile``. + +.. code-block:: python + + compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi") + + +Working with torch Tensors +-------------------------- + +As you may have noticed in the examples above, TVM FFI-compiled functions can +directly accept ``torch.Tensor`` objects (and other DLPack-compatible tensors) as inputs. +The resulting functions add minimal overhead, enabling faster eager invocations +thanks to the optimized calling path. + +Working with Streams +-------------------- + +In many cases, a CuTe kernel needs to run on a specific CUDA stream. +|DSL| provides two ways to work with streams through TVM FFI. +The first is to pass the stream explicitly as an argument. +The following example demonstrates this approach; the function accepts ``torch.cuda.Stream``, +``CUstream`` or any stream class that implements the CUDA stream protocol. + +.. code-block:: python + + import cutlass.cute as cute + import torch + from cuda.bindings.driver import CUstream + + @cute.kernel + def device_add_one(a: cute.Tensor, b: cute.Tensor): + threads_per_block = 128 + cta_x_, _, _ = cute.arch.block_idx() + tid_x, _, _ = cute.arch.thread_idx() + tid = cta_x_ * threads_per_block + tid_x + if tid < a.shape[0]: + b[tid] = a[tid] + 1.0 + + @cute.jit + def add_one_with_stream(a: cute.Tensor, b: cute.Tensor, stream: CUstream): + n = a.shape[0] + threads_per_block = 128 + blocks = (n + threads_per_block - 1) // threads_per_block + device_add_one(a, b).launch( + grid=(blocks, 1, 1), + block=(threads_per_block, 1, 1), + stream=stream, + ) + + def example_add_one_with_stream(): + n = cute.sym_int() + a_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,)) + b_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,)) + # Fake stream is a placeholder for stream argument + stream = cute.runtime.make_fake_stream() + compiled_add_one = cute.compile( + add_one_with_stream, a_cute, b_cute, stream, options="--enable-tvm-ffi" + ) + a_torch = torch.arange(10, dtype=torch.float32, device="cuda") + b_torch = torch.empty(10, dtype=torch.float32, device="cuda") + torch_stream = torch.cuda.current_stream() + compiled_add_one(a_torch, b_torch, torch_stream) + torch_stream.synchronize() + print("result of b_torch after compiled_add_one(a_torch, b_torch, torch_stream)") + print(b_torch) + +The second option is to rely on the environment-stream flag. +Pass ``use_tvm_ffi_env_stream=True`` to ``make_fake_stream`` to mark the argument as an +environment stream so it no longer has to be provided explicitly. +TVM FFI will reuse its environment stream, synchronizing it with ``torch.cuda.current_stream()`` +before each call. The example below shows this flow: + +.. code-block:: python + + def example_add_one_with_env_stream(): + n = cute.sym_int() + a_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,)) + b_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,)) + # Fake stream is a placeholder for stream argument + # we will use TVM FFI environment stream + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + compiled_add_one = cute.compile( + add_one_with_stream, a_cute, b_cute, stream, options="--enable-tvm-ffi" + ) + a_torch = torch.arange(10, dtype=torch.float32, device="cuda") + b_torch = torch.empty(10, dtype=torch.float32, device="cuda") + torch_stream = torch.cuda.current_stream() + with torch.cuda.stream(torch_stream): + # no need to pass in the stream explicitly, env stream will be synced + # to torch.cuda.current_stream() before the function call. + compiled_add_one(a_torch, b_torch) + torch_stream.synchronize() + print("result of b_torch after compiled_add_one(a_torch, b_torch)") + print(b_torch) + +Using the environment-stream flag both speeds up calls and simplifies integration +with frameworks such as PyTorch, since no explicit stream parameter is required. + +Supported types +--------------- + +The TVM FFI function supports the following |DSL|-specific types as arguments: + +- ``cute.Tensor`` +- ``cutlass.Boolean``, ``cutlass.Int8``, ``cutlass.Int16``, ``cutlass.Int32``, ``cutlass.Int64``, ``cutlass.Uint8``, ``cutlass.Uint16``, ``cutlass.Uint32``, ``cutlass.Uint64``, ``cutlass.Float32``, ``cutlass.Float64`` +- ``cute.Shape``, ``cute.Stride``, ``cute.Coord``, ``cute.Tile``, ``cute.IntTuple`` + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Compile-time type + - Call-time type + * - ``cute.Pointer`` + - ``ctypes.c_void_p`` or a class that implements ``__tvm_ffi_opaque_ptr__`` protocol. + * - ``cute.runtime.FakeTensor`` + - ``torch.Tensor`` and other DLPack-compatible tensors. + * - Scalar types (e.g. ``cutlass.Boolean``, ``cutlass.Int32``) + - Python scalars (e.g. True, 123). + * - CuTe algebra types (e.g. ``cute.Shape``, ``cute.Stride``) + - ``tvm_ffi.Shape`` or python tuple of ints. + * - CUDA stream ``cuda.CUstream`` + - A stream class that implements the CUDA stream protocol (e.g. ``torch.cuda.Stream``, ``cuda.CUstream``). + +Error handling +-------------- + +TVM FFI functions will enable validation of arguments to make sure they match the expected type +and value constraints declared by the user. These checks are compiled into the function, run very fast, +and have no observable overhead during function invocation. Each of those errors will translate +into a proper Python exception that can be caught and handled. The example below shows some +example error cases that can be checked: + +.. code-block:: python + + def example_constraint_checks(): + n = cute.sym_int(divisibility=16) + # assume align to 16 bytes (4 int32), both should share same shape variable n + a_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,), assumed_align=16) + b_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,), assumed_align=16) + compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi") + a = torch.zeros(128, dtype=torch.float32, device="cuda") + b = torch.zeros(128, dtype=torch.float32, device="cuda") + + try: + # raises type mismatch error because we expect a and b to be float32 + compiled_add_one(a, 1) + except TypeError as e: + # Mismatched type on argument #1 when calling: + # `add_one(a: Tensor([n0], float32), b: Tensor([n0], float32))`, + # expected Tensor + print(f"TypeError: {e}") + + try: + # raises shape mismatch error because we expect both a and b have shap [n] + compiled_add_one(a, b[:126]) + except ValueError as e: + # Mismatched b.shape[0] on argument #1 when calling: + # `add_one(a: Tensor([n0], float32), b: Tensor([n0], float32))`, + # symbolic constraint violated + print(f"ValueError: {e}") + + try: + # triggers divisibility mismatch error because 126 is not divisible by 16 + compiled_add_one(a[:126], b[:126]) + except ValueError as e: + # Invalid a.shape[0] on argument #0 when calling: + # `add_one(a: Tensor([n0], float32), b: Tensor([n0], float32)`, + # expected to be divisible by 16 + print(f"ValueError: {e}") + + try: + a = torch.zeros(129, dtype=torch.float32, device="cuda") + b = torch.zeros(129, dtype=torch.float32, device="cuda") + # triggers data alignment mismatch error because x and y are not aligned to 16 bytes + compiled_add_one(a[1:], b[1:]) + except ValueError as e: + # raises: Misaligned Tensor data on argument #0 when calling: + # `add_one(a: Tensor([n0], float32), b: Tensor([n0], float32)`, + # expected data alignment=16 bytes + print(f"ValueError: {e}") + +Any CUDA errors encountered will also be automatically converted into Python exceptions by the TVM FFI function. + +.. code-block:: python + + @cute.jit + def add_one_invalid_launch(a: cute.Tensor, b: cute.Tensor): + # Intentionally exceed the maximum block dimension (1024 threads) so the + # CUDA runtime reports an invalid configuration error. + device_add_one(a, b).launch(grid=(1, 1, 1), block=(4096, 1, 1)) + + def example_error_cuda_error(): + a_torch = torch.zeros((10,), dtype=torch.float32, device="cuda") + b_torch = torch.zeros((10,), dtype=torch.float32, device="cuda") + + a_cute = cute.runtime.from_dlpack(a_torch, enable_tvm_ffi=True) + b_cute = cute.runtime.from_dlpack(b_torch, enable_tvm_ffi=True) + compiled_add_one_invalid_launch = cute.compile( + add_one_invalid_launch, a_cute, b_cute, options="--enable-tvm-ffi" + ) + + try: + compiled_add_one_invalid_launch(a_torch, b_torch) + except RuntimeError as e: + # raises RuntimeError: CUDA Error: cudaErrorInvalidValue + print(f"RuntimeError: {e}") + + +Working with Devices +-------------------- +TVM FFI-compiled functions naturally work across GPU devices. +The device index of the first input GPU tensor determines the kernel's device context. +The TVM FFI function calls ``cudaSetDevice`` to set the correct device +before launching the kernel based on that tensor's device index. +For advanced scenarios that pass raw pointers instead of tensors, you should call +``cudaSetDevice`` explicitly through the CUDA Python API. + +Exporting Compiled Module +------------------------- + +The TVM FFI function supports exporting the compiled module to an object file +for further use. For example: + +.. code-block:: python + + import subprocess + import cutlass.cute as cute + + def example_add_one_export(): + n = cute.sym_int() + a_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,)) + b_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,)) + # compile the kernel with "--enable-tvm-ffi" option and example input tensors + compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi") + # export the compiled module to object file + compiled_add_one.export_to_c("./add_one.o", function_name="add_one") + # obtain necessary runtime libs for loading the shared library + runtime_libs = cute.runtime.find_runtime_libraries(enable_tvm_ffi=True) + # compile the object file to a shared library + cmd = ["gcc", "-shared", "-o", "./add_one.so", "./add_one.o", *runtime_libs] + print(cmd) + subprocess.run(cmd, check=True) + print(f"Successfully created shared library: ./add_one.so") + +Then you can load back the exported module and use it in different ways: + +.. code-block:: python + + import torch + from cutlass import cute + + def example_load_module_add_one(): + mod = cute.runtime.load_module("./add_one.so") + a_torch = torch.arange(10, dtype=torch.float32, device="cuda") + b_torch = torch.empty(10, dtype=torch.float32, device="cuda") + mod.add_one(a_torch, b_torch) + print("result of b_torch after mod.add_one(a_torch, b_torch)") + print(b_torch) + + +The exported object file exposes the function symbol ``__tvm_ffi_add_one`` that is +compatible with TVM FFI and can be used in various frameworks and programming languages. +You can either build a shared library and load it back, or link the object file directly +into your application and invoke the function via the ``InvokeExternC`` mechanism in TVM FFI. +For more information, see the `quick start guide `_ +in the official documentation. + +When you build your own libraries, make sure you link against the necessary runtime libraries. +You can use ``cute.runtime.find_runtime_libraries(enable_tvm_ffi=True)`` to get the path to these libraries. +``cute.runtime.load_module`` will load these libraries automatically before loading +an exported module. You can also manually load these libraries in advanced use cases. diff --git a/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.rst b/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.rst index 1f4bcb25..5f7e7598 100644 --- a/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.rst +++ b/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.rst @@ -62,6 +62,10 @@ You can provide additional compilation options as a string when calling ``cute.c - The GPU architecture to compile for. - "" - str + * - ``enable-tvm-ffi`` + - Enable Apache TVM FFI. + - False + - bool You can use the following code to specify compilation options: diff --git a/media/docs/pythonDSL/limitations.rst b/media/docs/pythonDSL/limitations.rst index 7b42c503..23673195 100644 --- a/media/docs/pythonDSL/limitations.rst +++ b/media/docs/pythonDSL/limitations.rst @@ -17,9 +17,7 @@ the DSL. Notable unsupported features ---------------------------- -- GeForce RTX 50 Series support - Programmatic Dependent Launch (PDL) -- narrow-precision data type support, including related tensor core instructions - convolutions - full support for ahead of time compilation - preferred clusters diff --git a/python/CuTeDSL/cutlass/base_dsl/_mlir_helpers/arith.py b/python/CuTeDSL/cutlass/base_dsl/_mlir_helpers/arith.py index 9afb8d05..44a83c7c 100644 --- a/python/CuTeDSL/cutlass/base_dsl/_mlir_helpers/arith.py +++ b/python/CuTeDSL/cutlass/base_dsl/_mlir_helpers/arith.py @@ -689,7 +689,21 @@ def _min(lhs, rhs, *, loc=None, ip=None): if not is_dynamic_expression(rhs): rhs = arith.constant(lhs.type, rhs, loc=loc, ip=ip) - if arith._is_integer_like_type(lhs.type): + # Handle vector types + if isinstance(lhs.type, T.VectorType): + elem_type = lhs.type.element_type + if arith._is_integer_like_type(elem_type): + assert hasattr( + lhs, "signed" + ), "Should have attribute `signed`, must be a bug" + if lhs.signed != False: + return arith.minsi(lhs, rhs, loc=loc, ip=ip) + else: + return arith.minui(lhs, rhs, loc=loc, ip=ip) + else: + return arith.minimumf(lhs, rhs, loc=loc, ip=ip) + elif arith._is_integer_like_type(lhs.type): + assert hasattr(lhs, "signed"), "Should have attribute `signed`, must be a bug" if lhs.signed != False: return arith.minsi(lhs, rhs, loc=loc, ip=ip) else: @@ -705,7 +719,6 @@ def _max(lhs, rhs, *, loc=None, ip=None): Assuming the operands have the same type """ from ..dsl import is_dynamic_expression - if not is_dynamic_expression(lhs): if not is_dynamic_expression(rhs): return max(lhs, rhs) @@ -714,8 +727,21 @@ def _max(lhs, rhs, *, loc=None, ip=None): else: if not is_dynamic_expression(rhs): rhs = arith.constant(lhs.type, rhs, loc=loc, ip=ip) - - if arith._is_integer_like_type(lhs.type): + # Handle vector types + if isinstance(lhs.type, T.VectorType): + elem_type = lhs.type.element_type + if isinstance(elem_type, ir.IntegerType): + assert hasattr( + lhs, "signed" + ), "Should have attribute `signed`, must be a bug" + if lhs.signed != False: + return arith.maxsi(lhs, rhs, loc=loc, ip=ip) + else: + return arith.maxui(lhs, rhs, loc=loc, ip=ip) + else: + return arith.maximumf(lhs, rhs, loc=loc, ip=ip) + elif arith._is_integer_like_type(lhs.type): + assert hasattr(lhs, "signed"), "Should have attribute `signed`, must be a bug" if lhs.signed != False: return arith.maxsi(lhs, rhs, loc=loc, ip=ip) else: diff --git a/python/CuTeDSL/cutlass/base_dsl/arch.py b/python/CuTeDSL/cutlass/base_dsl/arch.py index 78b9bb67..86876172 100644 --- a/python/CuTeDSL/cutlass/base_dsl/arch.py +++ b/python/CuTeDSL/cutlass/base_dsl/arch.py @@ -11,7 +11,7 @@ from enum import Enum import re -from typing import Callable, List +from typing import Callable, List, Tuple class Arch(Enum): @@ -49,16 +49,47 @@ class Arch(Enum): @classmethod def _missing_(cls, value): - """Support creating Arch enum from (major, minor, suffix) tuple""" - if not isinstance(value, tuple): - raise ValueError(f"invalid arguments for Arch: {value}") - major, minor, suffix = None, None, None - if len(value) == 2: + if isinstance(value, tuple) and len(value) == 2: + # Support creating Arch enum from (major, minor) tuple + # Arch(major, minor) is equivalent to Arch(major, minor, "") major, minor, suffix = *value, "" + return cls((major, minor, suffix)) else: raise ValueError(f"invalid arguments for Arch: {value}") + - return cls(major, minor, suffix) + # attributes to get arch list of specific families + @classmethod + def AmpereArchs(cls) -> Tuple["Arch"]: + return (Arch.sm_80, Arch.sm_86, Arch.sm_87) + + @classmethod + def AdaArchs(cls) -> Tuple["Arch"]: + return (Arch.sm_89,) + + @classmethod + def HopperArchs(cls) -> Tuple["Arch"]: + return (Arch.sm_90, Arch.sm_90a) + + @classmethod + def BlackwellArchs(cls) -> Tuple["Arch"]: + return ( + Arch.sm_100, + Arch.sm_100a, + Arch.sm_100f, + Arch.sm_101, + Arch.sm_101a, + Arch.sm_101f, + Arch.sm_110, + Arch.sm_110a, + Arch.sm_110f, + Arch.sm_120, + Arch.sm_120a, + Arch.sm_120f, + Arch.sm_121, + Arch.sm_121a, + Arch.sm_121f, + ) def __repr__(self): return self.__str__() diff --git a/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py b/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py index a4531914..16dc74ca 100644 --- a/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py +++ b/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py @@ -526,6 +526,10 @@ class DSLPreprocessor(ast.NodeTransformer): self.generic_visit(node) self.loop_nest_level -= 1 + def visit_FunctionDef(self, node): + # Stop at nested function def + return + checker = EarlyExitChecker(kind) checker.generic_visit(tree) if not checker.has_early_exit: @@ -1325,6 +1329,7 @@ class DSLPreprocessor(ast.NodeTransformer): node, ) elif func.id in ["min", "max"]: + self.import_top_module = True return ast.copy_location( ast.Call( func=self._create_module_attribute( @@ -1334,7 +1339,7 @@ class DSLPreprocessor(ast.NodeTransformer): lineno=node.lineno, col_offset=node.col_offset, ), - args=[node.args[0], node.args[1]], + args=node.args, keywords=[], ), node, diff --git a/python/CuTeDSL/cutlass/base_dsl/common.py b/python/CuTeDSL/cutlass/base_dsl/common.py index 7d36e01d..6bc5fca1 100644 --- a/python/CuTeDSL/cutlass/base_dsl/common.py +++ b/python/CuTeDSL/cutlass/base_dsl/common.py @@ -134,7 +134,7 @@ def _get_friendly_cuda_error_message(error_code, error_name): # Add target architecture info target_arch = os.getenv("CUTE_DSL_ARCH", "unknown") - error_messages = { + additional_info = { "CUDA_ERROR_INVALID_SOURCE": ( f"{Colors.RED}❌ Failed to load CUDA kernel - likely architecture mismatch.{Colors.RESET}\n\n" ), @@ -157,6 +157,9 @@ def _get_friendly_cuda_error_message(error_code, error_name): f"{Colors.RED}⚠️ Invalid parameter passed to CUDA operation.{Colors.RESET}\n\n" f"{Colors.YELLOW}This is likely a bug - please report it with:{Colors.RESET}" ), + "CUDA_ERROR_INVALID_CLUSTER_SIZE": ( + f"{Colors.RED}❌ Invalid cluster size.{Colors.RESET}\n\n" + ), } error_suggestions = { @@ -194,12 +197,12 @@ def _get_friendly_cuda_error_message(error_code, error_name): ), } - message = error_messages.get( - error_name, f"{Colors.RED}Unknown CUDA error{Colors.RESET}" - ) + message = f"{error_name} (error code: {error_code}) \n" \ + f"{additional_info.get(error_name, '')} \n\n{Colors.RESET}" # Add debug information debug_info = f"\n- {Colors.BOLD}Error name: {error_name}\n" + debug_info += f"- Error code: {error_code}\n" debug_info += f"- CUDA_TOOLKIT_PATH: {os.getenv('CUDA_TOOLKIT_PATH', 'not set')}\n" debug_info += ( f"- Target SM ARCH: {os.getenv('CUTE_DSL_ARCH', 'not set')}{Colors.RESET}\n" diff --git a/python/CuTeDSL/cutlass/base_dsl/compiler.py b/python/CuTeDSL/cutlass/base_dsl/compiler.py index 9aee3acc..3b47c578 100644 --- a/python/CuTeDSL/cutlass/base_dsl/compiler.py +++ b/python/CuTeDSL/cutlass/base_dsl/compiler.py @@ -15,7 +15,7 @@ and executes it using MLIR's ExecutionEngine. """ -from typing import Sequence, Optional, Tuple +from typing import Sequence, Optional, Tuple, Callable import os import sys import inspect @@ -99,6 +99,8 @@ class Compiler: self.execution_engine = execution_engine # Flag to track if CUDA dependencies have been checked once in this process self._cuda_dependencies_checked = False + # Post-compile hook to run on Module + self._post_compile_hook: Optional[Callable[[ir.Module], None]] = None def _process_error(self, error_msg: str) -> Tuple[Optional[str], Optional[str]]: """Process error message to extract NVVM error and IR context""" @@ -158,11 +160,13 @@ class Compiler: ) from e raise e + if self._post_compile_hook: + self._post_compile_hook(module) + def jit(self, module, opt_level: int = 2, shared_libs: Sequence[str] = ()): """Wraps the module in a JIT execution engine.""" # Check CUDA driver and GPU dependencies before JIT execution (once per process) self._check_cuda_dependencies_once(shared_libs) - # If pre-checks passed, attempt to create ExecutionEngine # Any failures at this point are likely non-CUDA related return self.execution_engine.ExecutionEngine( @@ -216,6 +220,23 @@ class Compiler: ) +class PostCompileHookContext: + """Context manager for post-compile hook for a compiler.""" + + def __init__(self, compiler: Compiler, hook: Callable[[ir.Module], None]): + self.compiler = compiler + self.hook = hook + self.prev_post_compile_hook = None + + def __enter__(self): + self.prev_post_compile_hook = self.compiler._post_compile_hook + self.compiler._post_compile_hook = self.hook + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.compiler._post_compile_hook = self.prev_post_compile_hook + + class CompileOption: """ Base class for compile options. @@ -276,6 +297,11 @@ class BooleanBasedFileDumpOption(CompileOption): return "" +class EmptyCompileOption(CompileOption): + def serialize(self): + return "" + + class OptLevel(CompileOption): option_name = "opt-level" @@ -329,6 +355,10 @@ class GPUArch(StringCompileOption): self._value = value +class EnableTVMFFI(EmptyCompileOption): + option_name = "enable-tvm-ffi" + + class CompileOptions: """ This class encapsulates compilation options to configure the JIT compilation. @@ -349,6 +379,7 @@ class CompileOptions: KeepPTX: KeepPTX(False), GPUArch: GPUArch(""), LinkLibraries: LinkLibraries(""), + EnableTVMFFI: EnableTVMFFI(False), } if options is not None: @@ -376,15 +407,28 @@ class CompileOptions: self.options[EnableAssertions].value = True if envar.lineinfo: self.options[GenerateLineInfo].value = True + if envar.enable_tvm_ffi: + self.options[EnableTVMFFI].value = True # Update the dump path if the option is set + arch = ( + envar.arch + if self.options[GPUArch].value == "" + else self.options[GPUArch].value + ) if self.options[KeepPTX].value: self.options[KeepPTX].dump_path = os.path.join( - envar.dump_dir, f"{function_name}.ptx" + envar.dump_dir, f"{function_name}" + ) + self.options[KeepPTX].full_ptx_path = os.path.join( + envar.dump_dir, f"{function_name}.{arch}.ptx" ) if self.options[KeepCUBIN].value: self.options[KeepCUBIN].dump_path = os.path.join( - envar.dump_dir, f"{function_name}.cubin" + envar.dump_dir, f"{function_name}" + ) + self.options[KeepCUBIN].full_cubin_path = os.path.join( + envar.dump_dir, f"{function_name}.{arch}.cubin" ) @property @@ -399,12 +443,38 @@ class CompileOptions: def dump_ptx_path(self) -> str | None: return self.options[KeepPTX].dump_path if self.options[KeepPTX].value else None + @property + def full_ptx_path(self) -> str | None: + return ( + self.options[KeepPTX].full_ptx_path if self.options[KeepPTX].value else None + ) + @property def dump_cubin_path(self) -> str | None: return ( self.options[KeepCUBIN].dump_path if self.options[KeepCUBIN].value else None ) + @property + def full_cubin_path(self) -> str | None: + return ( + self.options[KeepCUBIN].full_cubin_path + if self.options[KeepCUBIN].value + else None + ) + + @property + def enable_tvm_ffi(self) -> bool: + ret = self.options[EnableTVMFFI].value + if ret: + try: + import tvm_ffi + except ModuleNotFoundError: + raise DSLRuntimeError( + "TVM FFI is not installed, please install it via `pip install apache-tvm-ffi`" + ) + return ret + def to_str(self) -> str: """ Generate a string representation of all compilation options @@ -433,6 +503,7 @@ def _parse_compile_options_from_str(options: str) -> CompileOptions: "keep_cubin": KeepCUBIN, "keep_ptx": KeepPTX, "gpu_arch": GPUArch, + "enable_tvm_ffi": EnableTVMFFI, } return mapping[option_str] @@ -448,6 +519,7 @@ def _parse_compile_options_from_str(options: str) -> CompileOptions: parser.add_argument("--keep-ptx", action="store_true", default=False) parser.add_argument("--ptxas-options", type=str, default="") parser.add_argument("--gpu-arch", type=str, default="") + parser.add_argument("--enable-tvm-ffi", action="store_true", default=False) compile_options = CompileOptions() try: # Use shlex to properly handle options with spaces @@ -553,6 +625,9 @@ class CompileCallable: # process compile options, extract the options and remove them from the kwargs options = kwargs.pop("options", None) + if isinstance(options, str) and len(options) == 0: + options = None + if options is not None and isinstance(options, str): compile_options = _parse_compile_options_from_str(options) else: diff --git a/python/CuTeDSL/cutlass/base_dsl/dsl.py b/python/CuTeDSL/cutlass/base_dsl/dsl.py index 4cb3c111..e9e8662f 100644 --- a/python/CuTeDSL/cutlass/base_dsl/dsl.py +++ b/python/CuTeDSL/cutlass/base_dsl/dsl.py @@ -27,6 +27,7 @@ import re import inspect import argparse import hashlib +import weakref from functools import lru_cache, wraps from collections import namedtuple, OrderedDict from abc import ABC, abstractmethod @@ -52,7 +53,7 @@ from .runtime.jit_arg_adapters import is_argument_constexpr, JitArgAdapterRegist from .ast_preprocessor import DSLPreprocessor from .common import * -from .typing import get_c_pointers, get_mlir_types, Integer +from .typing import get_c_pointers, get_mlir_types, Integer, arg_compatible_with_tvm_ffi from .arch import Arch # ============================================================================= @@ -345,6 +346,7 @@ class BaseDSL: if preprocess: self.preprocessor = DSLPreprocessor(dsl_package_name) + log().info(f"Initializing {name} DSL") log().debug(f"Logger initialized for {self.name}") @@ -756,6 +758,7 @@ class BaseDSL: args_spec: inspect.FullArgSpec, *, is_host=True, + compile_only=False, ): """Generate JIT function arguments.""" @@ -806,7 +809,14 @@ class BaseDSL: jit_adapted_args.append(arg) if is_host: - jit_exec_arg.extend(get_c_pointers(arg)) + if self.envar.enable_tvm_ffi: + if not arg_compatible_with_tvm_ffi(arg): + raise DSLRuntimeError( + f"Argument #{i + 1} ({arg_name}) is not a TVM FFI argument." + ) + jit_exec_arg.extend([arg]) + else: + jit_exec_arg.extend(get_c_pointers(arg)) jit_arg_type.extend(get_mlir_types(arg)) else: dyn_vals = extract_mlir_values(arg) @@ -814,7 +824,10 @@ class BaseDSL: jit_arg_type.extend([v.type for v in dyn_vals]) if not jit_arg_type or not jit_exec_arg: - if (is_host and hasattr(arg, "__c_pointers__")) or ( + # when it is compile only, we don't have to prepare the executable arguments. + if ( + is_host and (compile_only or hasattr(arg, "__c_pointers__")) + ) or ( not is_host and hasattr(arg, "__extract_mlir_values__") and hasattr(arg, "__new_from_mlir_values__") @@ -845,20 +858,32 @@ class BaseDSL: return jit_exec_args, jit_arg_types, jit_arg_attrs, jit_adapted_args def generate_mlir_function_types( - self, func, function_name, input_args, kwargs, args_spec: inspect.FullArgSpec + self, + func, + function_name, + input_args, + kwargs, + args_spec: inspect.FullArgSpec, + compile_only=False, ): """Convert input arguments to MLIR function signature also convert numpy arrays to memref.""" exe_args, types, attrs, adapted_args = self._generate_jit_func_args( - func, function_name, input_args, kwargs, args_spec, is_host=True + func, + function_name, + input_args, + kwargs, + args_spec, + is_host=True, + compile_only=compile_only, ) log().debug("Execution Arguments: %s", ", ".join(map(str, exe_args))) log().debug("Types: %s", ", ".join(map(str, types))) - assert len(exe_args) == len(types), ( - "expects the same number of arguments and function parameters" - ) + assert ( + compile_only or self.envar.enable_tvm_ffi or len(exe_args) == len(types) + ), "expects the same number of arguments and function parameters" return exe_args, types, adapted_args @@ -867,6 +892,7 @@ class BaseDSL: cluster: list = None grid: list = field(default_factory=lambda: [1, 1, 1]) block: list = field(default_factory=lambda: [1, 1, 1]) + max_number_threads: list = field(default_factory=lambda: [0, 0, 0]) smem: int = None async_deps: list = field(default_factory=list) has_cluster: bool = False @@ -906,6 +932,13 @@ class BaseDSL: elif len(self.cluster) != 3: raise DSLRuntimeError(f"Expect 3d cluster!") + def has_max_number_threads(self): + """Check if max_number_threads is given by user""" + return all( + value == 0 if not is_dynamic_expression(value) else False + for value in self.max_number_threads + ) + def diagnostic(self): """Check command line parameters and enables diagnostic""" # Check command line arguments "-diagnostic" @@ -1029,7 +1062,6 @@ class BaseDSL: pipeline = re.sub(r"{.+}", opt_str, pipeline) else: pipeline = pipeline.rstrip(")") + f"{{{opt_str}}})" - log().debug(f"Using pipeline = {pipeline}") return pipeline def get_shared_libs(self) -> list: @@ -1101,6 +1133,18 @@ class BaseDSL: return module + def get_return_types(self) -> List[ir.Type]: + """ + Get the return types of the host function. + """ + return [] + + def generate_default_return_values(self, ip=None) -> List[ir.Value]: + """ + Generate the default return values of the host function. + """ + return [] + def generate_original_ir( self, ir, @@ -1115,27 +1159,35 @@ class BaseDSL: frame=None, ): def build_ir_module(): - module = ir.Module.create(loc=self.get_location(frame)) + loc = self.get_location(frame) + module = ir.Module.create(loc=loc) unit_attr = ir.UnitAttr.get() module.operation.attributes["gpu.container_module"] = unit_attr with ir.InsertionPoint(module.body): # Always generate gpu module. It's canonicalized by the compiler when it's not used. - self._build_gpu_module(gpu_module_attrs, loc=self.get_location(frame)) + self._build_gpu_module(gpu_module_attrs, loc=loc) + ret_types = self.get_return_types() fop = func.FuncOp( - function_name, (func_types, []), loc=self.get_location(frame) + function_name, (func_types, ret_types), loc=loc ) fop.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get() log().debug("Generated Function OP [%s]", fop) - with ir.InsertionPoint(fop.add_entry_block()): + # Attach per-argument source locations if supported by the FuncOp binding. + arg_locs = [loc for _ in range(len(func_types))] + entry_block = fop.add_entry_block(arg_locs=arg_locs) + with ir.InsertionPoint(entry_block): ir_args, ir_kwargs = self.generate_execution_arguments( args, kwargs, fop, args_spec ) # Call user function body try: result = funcBody(*ir_args, **ir_kwargs) - func.ReturnOp([], loc=self.get_location(frame)) + default_ret_values = self.generate_default_return_values( + ir.InsertionPoint.current + ) + func.ReturnOp(default_ret_values, loc=loc) except NameError as name_error: raise DSLRuntimeError( f"💥💥💥 Error during runtime code generation for function `{funcBody.__name__}` 💥💥💥", @@ -1165,6 +1217,10 @@ class BaseDSL: args_spec, no_cache, func_type=JitCompiledFunction, + *, + dynamic_args=None, + dynamic_kwargs=None, + original_function_name=None, ): # If `gpu-arch` is set by compile_options, use it. Otherwise, use the arch from the environment variable. compile_gpu_arch = ( @@ -1178,6 +1234,7 @@ class BaseDSL: pipeline = self.preprocess_pipeline( self._get_pipeline(pipeline), compile_gpu_arch ) + log().debug(f"Using pipeline = {pipeline}") shared_libs = self.get_shared_libs() profiler = timer(enable=self.envar.jit_time_profiling) if ( @@ -1212,7 +1269,7 @@ class BaseDSL: ) capi_func = profiler(engine.lookup)(function_name) if engine else None - fn = JitCompiledFunction( + fn = func_type( module, engine, capi_func, @@ -1221,12 +1278,15 @@ class BaseDSL: self.kernel_info, jit_time_profiling=self.envar.jit_time_profiling, jit_function_artifacts=JitFunctionArtifacts( - PTX=self.compile_options.dump_ptx_path, - CUBIN=self.compile_options.dump_cubin_path, + PTX=self.compile_options.full_ptx_path, + CUBIN=self.compile_options.full_cubin_path, MLIR=(self.dump_mlir_path if self.envar.keep_ir else None), ), ) + # set dynamic arguments if the jit_function is a JitCompiledFunction for AOT generation. + fn.set_dynamic_args(dynamic_args, dynamic_kwargs) + if not no_cache: # module stored in cache is compiled. self.jit_cache[module_hash] = fn @@ -1243,6 +1303,34 @@ class BaseDSL: # reset the compile options after the compilation is done. self.compile_options = CompileOptions() + def extract_dynamic_args(self, funcBody, args, kwargs, args_spec): + """This function is used to extract the original dynamic arguments for AOT C header generation. + The dynamic argument is the argument which is not marked as `Constexpr` in the function signature. + """ + dynamic_args = [] + dynamic_kwargs = OrderedDict() + for i, arg in enumerate(args): + if not is_argument_constexpr( + arg, + args_spec.annotations.get(args_spec.args[i], None), + args_spec.args[i], + i, + funcBody, + ): + try: + dynamic_args.append(weakref.proxy(arg)) + except TypeError: + # If arg cannot be weakly referenced (e.g., int, float) + 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): + try: + dynamic_kwargs[k] = weakref.proxy(v) + except TypeError: + # If v cannot be weakly referenced (e.g., int, float) + dynamic_kwargs[k] = v + return dynamic_args, dynamic_kwargs + def generate_mlir( self, funcBody, @@ -1258,49 +1346,63 @@ class BaseDSL: frame=None, ): """Generate MLIR module and compile iself.T_provider.""" - with ir.Context(), ir.Location.unknown(): - # Convert input arguments to MLIR arguments - exe_args, func_types, adapted_args = self.generate_mlir_function_types( - funcBody, function_name, args, kwargs, args_spec - ) - - # Generate original ir module and its hash value. - module, module_hash, result = self.generate_original_ir( - ir, - func, - funcBody, - kwargs, - function_name, - func_types, - gpu_module_attrs, - args, - args_spec, - frame=frame, - ) - - # dryrun is used to only generate IR - if self.envar.dryrun: - return result - - if ( - no_cache - or module_hash not in self.jit_cache - or self.jit_cache[module_hash].capi_func is None - ): - # no cache or cache miss, do ir generation/compilation/jit engine - jit_function = self.compile_and_cache( - module, module_hash, function_name, pipeline, args_spec, no_cache + with ir.Context(), self.get_location(frame): + try: + # Convert input arguments to MLIR arguments + exe_args, func_types, adapted_args = self.generate_mlir_function_types( + funcBody, function_name, args, kwargs, args_spec, compile_only ) - else: - # cache hit - log().info( - "JIT cache hit IN-MEMORY function=[%s] module_hash=[%s]", + dynamic_args, dynamic_kwargs = self.extract_dynamic_args( + funcBody, args, kwargs, args_spec + ) + original_function_name = funcBody.__name__ + + # Generate original ir module and its hash value. + module, module_hash, result = self.generate_original_ir( + ir, + func, + funcBody, + kwargs, function_name, - module_hash, + func_types, + gpu_module_attrs, + args, + args_spec, + frame=frame, ) - jit_function = self.jit_cache[module_hash] - self.post_compilation_cleanup() + # dryrun is used to only generate IR + if self.envar.dryrun: + return result + + if ( + no_cache + or module_hash not in self.jit_cache + or self.jit_cache[module_hash].capi_func is None + ): + # no cache or cache miss, do ir generation/compilation/jit engine + jit_function = self.compile_and_cache( + module, + module_hash, + function_name, + pipeline, + args_spec, + no_cache, + dynamic_args=dynamic_args, + dynamic_kwargs=dynamic_kwargs, + original_function_name=original_function_name, + ) + else: + # cache hit + log().info( + "JIT cache hit IN-MEMORY function=[%s] module_hash=[%s]", + function_name, + module_hash, + ) + jit_function = self.jit_cache[module_hash] + + finally: + self.post_compilation_cleanup() # If compile_only is set, bypass execution return the jit_executor directly if compile_only: @@ -1569,13 +1671,13 @@ class BaseDSL: """ ret = None - with ir.Context(), ir.Location.unknown(): + with ir.Context(), self.get_location(): loc = self.get_location() module = ir.Module.create(loc=loc) unit_attr = ir.UnitAttr.get() module.operation.attributes["gpu.container_module"] = unit_attr with ir.InsertionPoint(module.body): - self._build_gpu_module({}) + self._build_gpu_module({}, loc=loc) ret, kernel_name = kernel_generator() log().debug( f"Kernel generator returned: ret={ret}, kernel_name={kernel_name}" diff --git a/python/CuTeDSL/cutlass/base_dsl/env_manager.py b/python/CuTeDSL/cutlass/base_dsl/env_manager.py index e15cc757..94873e94 100644 --- a/python/CuTeDSL/cutlass/base_dsl/env_manager.py +++ b/python/CuTeDSL/cutlass/base_dsl/env_manager.py @@ -245,8 +245,15 @@ def get_prefix_dsl_libs(prefix: str): ] libs_cand = find_libs_in_ancestors(start, target_libs, lib_folder_guesses) + + optional_libs_cand = find_libs_in_ancestors( + start, {"cuda_dialect_runtime"}, lib_folder_guesses + ) + if libs_cand: dsl_libs = ":".join(libs_cand) + if optional_libs_cand: + dsl_libs += ":" + ":".join(optional_libs_cand) return dsl_libs return None @@ -312,6 +319,7 @@ class EnvironmentVarManager(LogEnvironmentManager): - [DSL_NAME]_DISABLE_FILE_CACHING: Disable file caching (default: False) - [DSL_NAME]_FILE_CACHING_CAPACITY: Limits the number of the cache save/load files (default: 1000) - [DSL_NAME]_LIBS: Path to dependent shared libraries (default: None) + - [DSL_NAME]_ENABLE_TVM_FFI: Enable TVM-FFI or not (default: False) """ def __init__(self, prefix="DSL"): @@ -358,3 +366,5 @@ class EnvironmentVarManager(LogEnvironmentManager): # whether to enable assert in host and device code self.enable_assertions = get_bool_env_var(f"{prefix}_ENABLE_ASSERTIONS", False) + + self.enable_tvm_ffi = get_bool_env_var(f"{prefix}_ENABLE_TVM_FFI", False) diff --git a/python/CuTeDSL/cutlass/base_dsl/export/__init__.py b/python/CuTeDSL/cutlass/base_dsl/export/__init__.py new file mode 100644 index 00000000..37e7412e --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/export/__init__.py @@ -0,0 +1,20 @@ +# 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 +# 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 .c_header_generator import CHeaderGenerator +from .export import get_export_module, dump_to_object, export_to_c + +__all__ = [ + "CHeaderGenerator", + "get_export_module", + "dump_to_object", + "export_to_c", +] diff --git a/python/CuTeDSL/cutlass/base_dsl/export/c_header_generator.py b/python/CuTeDSL/cutlass/base_dsl/export/c_header_generator.py new file mode 100644 index 00000000..4a22525a --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/export/c_header_generator.py @@ -0,0 +1,383 @@ +# 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 +# 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 ( + Numeric, + NumericMeta, + Int8, + Int16, + Int32, + Int64, + Uint8, + Uint16, + Uint32, + Uint64, + Boolean, + Float16, + Float32, + Float64, + TFloat32, +) +from ..dsl import is_dynamic_expression +from ..common import DSLRuntimeError +from ..jit_executor import ExecutionArgs +from ..._mlir import ir + +from typing import Type, List, Any, Dict +from inspect import isclass +import cuda.bindings.driver as cuda + +# ============================================================================= +# C Header Generator for c/cpp AOT support +# ============================================================================= +cubin_suffix = "cubin" + + +class CHeaderGenerator: + """This class provides a Export C Header Generator for c/cpp AOT support.""" + + includes = """ +#pragma once + +#include +#include +#include + +""" + cuda_error_check = r"""_CUDA_ERROR_CHECK(err) { \ + if ((err) != CUDA_SUCCESS) { \ + const char *pErrStr = NULL; \ + if (cuGetErrorString(err, &pErrStr) != CUDA_SUCCESS) { \ + printf("Error: cuGetErrorString failed\n"); \ + exit(err); \ + } \ + const char *pErrName = NULL; \ + if (cuGetErrorName(err, &pErrName) != CUDA_SUCCESS) { \ + printf("Error: cuGetErrorName failed\n"); \ + exit(err); \ + } \ + printf("Got Cuda Error: %s %s\n", pErrName, pErrStr); \ + exit(err); \ + } \ +} +""" + + numeric_to_c_type = { + Int8: "int8_t ", + Int16: "int16_t ", + Int32: "int32_t ", + Int64: "int64_t ", + Uint8: "uint8_t ", + Uint16: "uint16_t ", + Uint32: "uint32_t ", + Uint64: "uint64_t ", + Boolean: "uint8_t ", + Float32: "float ", + TFloat32: "float ", + Float64: "double ", + Float16: "__half_raw ", + } + + def _count_dynamic_expression(self, arg): + """ + Count the number of dynamic values in the argument. + """ + if isinstance(arg, (list, tuple)): + return sum(self._count_dynamic_expression(x) for x in arg) + elif is_dynamic_expression(arg): + return 1 + return 0 + + def _generate_numeric_argument(self, arg_name: str, arg_type: Type[Numeric]): + """ + Generate the argument of the wrapper function. + """ + if arg_type in self.numeric_to_c_type: + return self.numeric_to_c_type[arg_type] + arg_name + raise DSLRuntimeError( + f"Unsupported argument type for c function argument generation: {arg_type}" + ) + + def _generate_check_cuda(self, dsl_name: str): + check_cuda = ( + f""" +// Macro to check for cuda errors. +#ifndef {dsl_name}_CUDA_ERROR_CHECK +#define {dsl_name}""" + + self.cuda_error_check + + f""" +#endif +""" + ) + + return check_cuda + + def _generate_kernel_metadata( + self, symbol_prefix: str, kernel_info: Dict[str, List], dsl_name: str + ): + """ + Generate the kernel metadata for the compiled function. + """ + function_declarations = [] + function_loads = [] + function_set_attributes = [] + function_non_portable_cluster_size_allowed = [] + for sym, attrs in kernel_info.items(): + function_declaration = f"CUfunction {sym};" + function_load = f'{dsl_name}_CUDA_ERROR_CHECK(cuModuleGetFunction(&metadata->{sym}, metadata->module, "{sym}"));' + function_declarations.append(function_declaration) + function_loads.append(function_load) + for attr in attrs: + function_set_attributes.append( + f"{dsl_name}_CUDA_ERROR_CHECK(cuFuncSetAttribute(metadata->{sym}, {attr.name}, {attr.value}));" + ) + function_non_portable_cluster_size_allowed.append( + f"{dsl_name}_CUDA_ERROR_CHECK(cuFuncSetAttribute(metadata->{sym}, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1));" + ) + function_declarations_str = "\n ".join(function_declarations) + kernel_metadata_struct = f""" +typedef struct {{ + CUmodule module; + {function_declarations_str} +}} {symbol_prefix}_Kernel_Metadata_t; +""" + function_loads_str = "\n ".join(function_loads) + function_set_attributes_str = "\n ".join(function_set_attributes) + function_non_portable_cluster_size_allowed_str = "\n ".join( + function_non_portable_cluster_size_allowed + ) + kernel_metadata_load = f""" +static inline void {symbol_prefix}_Kernel_Metadata_Load({symbol_prefix}_Kernel_Metadata_t *metadata) {{ + {dsl_name}_CUDA_ERROR_CHECK(cuModuleLoadData(&metadata->module, {symbol_prefix}_{cubin_suffix})); + {function_loads_str} + {function_set_attributes_str} + int driver_version; + {dsl_name}_CUDA_ERROR_CHECK(cuDriverGetVersion(&driver_version)); + if (driver_version >= 11080) {{ + {function_non_portable_cluster_size_allowed_str} + }} +}} +""" + kernel_metadata_unload = f""" +static inline void {symbol_prefix}_Kernel_Metadata_Unload({symbol_prefix}_Kernel_Metadata_t *metadata) {{ + {dsl_name}_CUDA_ERROR_CHECK(cuModuleUnload(metadata->module)); +}} +""" + + return kernel_metadata_struct + kernel_metadata_load + kernel_metadata_unload + + def _generate_arguments( + self, + symbol_prefix: str, + args_spec: ExecutionArgs, + args: List[Any], + kwargs: Dict[str, Any], + ): + """ + Generate the arguments of the wrapper function. + """ + arguments = [] + packed_args = [] + declarations = [] + # traverse the runtime args_spec and generate the arguments + rectified_args = args_spec.get_rectified_args(args, kwargs) + input_arg_names = args_spec.args_spec.args + args_spec.args_spec.kwonlyargs + for arg_name, arg in zip(input_arg_names, rectified_args): + arg_type = args_spec.args_spec.annotations.get(arg_name, None) + + # process optional argument + if arg is None: + continue + + # Generate basic numeric types + if isinstance(arg_type, NumericMeta): + arguments.append(self._generate_numeric_argument(arg_name, arg_type)) + packed_args.append("&" + arg_name) + elif isclass(arg_type) and issubclass(arg_type, cuda.CUstream): + arguments.append("CUstream " + arg_name) + packed_args.append("&" + arg_name) + else: + raise DSLRuntimeError( + f"Unsupported argument for c function argument generation: {arg} with type {arg_type}" + ) + + return arguments, packed_args, declarations + + def _generate_wrapper_function( + self, + symbol_prefix: str, + args_spec: ExecutionArgs, + function_name: str, + kernel_info: Dict[str, List], + dynamic_args: list, + dynamic_kwargs: dict, + ): + """ + Generate the wrapper function for the compiled function which is provided to users as the entry point. + It uses the `symbol_prefix` as the function name for identification. The host/device symbols are hidden under the bytecode. + """ + # 1. Get the name of the function wrapper + wrapper_function_name = f"{symbol_prefix}_wrapper" + capi_function_name = f"_mlir_{symbol_prefix}__mlir_ciface_{function_name}" + # 2. Generate the signature of the wrapper function + arguments, packed_args, declarations = self._generate_arguments( + symbol_prefix, args_spec, dynamic_args, dynamic_kwargs + ) + declarations = "\n".join(declarations) + + # 3. Generate the wrapper function + kernel_symbols = tuple(kernel_info.keys()) + kernel_symbols_str = ", ".join([f"&metadata->{sym}" for sym in kernel_symbols]) + function = ( + declarations + + f""" +#ifdef __cplusplus +extern "C" +#endif +void {capi_function_name}(void **args, int32_t num_args); + +static inline void {wrapper_function_name}({symbol_prefix}_Kernel_Metadata_t *metadata, {", ".join(arguments)}) {{ + void *args[{len(packed_args) + len(kernel_symbols)}] = {{ + {", ".join(packed_args)}, {kernel_symbols_str} + }}; + {capi_function_name}(args, {len(packed_args) + len(kernel_symbols)}); +}} +""" + ) + return function + + def _generate_binary_declaration(self, symbol_prefix: str): + """ + Generate the binary of the compiled function. + """ + + varname = symbol_prefix + "_" + cubin_suffix + binary = f""" +extern const unsigned char {varname}[]; +""" + return binary + + def __call__( + self, + symbol_prefix: str, + export_module: ir.Module, + args_spec: ExecutionArgs, + function_name: str, + kernel_info: Dict[str, List], + dynamic_args: list, + dynamic_kwargs: dict, + dsl_name: str, + ) -> str: + if len(kernel_info) > 0: + check_cuda = self._generate_check_cuda(dsl_name) + kernel_metadata = self._generate_kernel_metadata( + symbol_prefix, kernel_info, dsl_name + ) + binary = self._generate_binary_declaration(symbol_prefix) + else: + check_cuda = "" + kernel_metadata = "" + binary = "" + function = self._generate_wrapper_function( + symbol_prefix, + args_spec, + function_name, + kernel_info, + dynamic_args, + dynamic_kwargs, + ) + header = self.includes + check_cuda + binary + kernel_metadata + function + return header + + +# ============================================================================= +# Below is the example output of a complete header file. +# ============================================================================= +# #pragma once + +# #include +# #include +# #include + + +# // Macro to check for cuda errors. +# #ifndef CUTE_DSL_CUDA_ERROR_CHECK +# #define CUTE_DSL_CUDA_ERROR_CHECK(err) { \ +# if ((err) != CUDA_SUCCESS) { \ +# const char *pErrStr = NULL; \ +# if (cuGetErrorString(err, &pErrStr) != CUDA_SUCCESS) { \ +# printf("Error: cuGetErrorString failed\n"); \ +# exit(err); \ +# } \ +# const char *pErrName = NULL; \ +# if (cuGetErrorName(err, &pErrName) != CUDA_SUCCESS) { \ +# printf("Error: cuGetErrorName failed\n"); \ +# exit(err); \ +# } \ +# printf("Got Cuda Error: %s %s\n", pErrName, pErrStr); \ +# exit(err); \ +# } \ +# } +# #endif + +# extern const unsigned char gemm_cubin[]; + +# typedef struct { +# CUmodule module; +# CUfunction kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0; +# } gemm_Kernel_Metadata_t; + +# static inline void gemm_Kernel_Metadata_Load(gemm_Kernel_Metadata_t *metadata) { +# CUTE_DSL_CUDA_ERROR_CHECK(cuModuleLoadData(&metadata->module, gemm_cubin)); +# CUTE_DSL_CUDA_ERROR_CHECK(cuModuleGetFunction(&metadata->kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0, metadata->module, "kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0")); + +# int driver_version; +# CUTE_DSL_CUDA_ERROR_CHECK(cuDriverGetVersion(&driver_version)); +# if (driver_version >= 11080) { +# CUTE_DSL_CUDA_ERROR_CHECK(cuFuncSetAttribute(metadata->kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1)); +# } +# } + +# static inline void gemm_Kernel_Metadata_Unload(gemm_Kernel_Metadata_t *metadata) { +# CUTE_DSL_CUDA_ERROR_CHECK(cuModuleUnload(metadata->module)); +# } + +# typedef struct { +# void *data; +# int32_t dynamic_shapes[2]; +# int64_t dynamic_strides[1]; +# } gemm_Tensor_a_t; + + +# typedef struct { +# void *data; +# int32_t dynamic_shapes[2]; +# int64_t dynamic_strides[1]; +# } gemm_Tensor_b_t; + + +# typedef struct { +# void *data; +# int32_t dynamic_shapes[2]; +# int64_t dynamic_strides[1]; +# } gemm_Tensor_c_t; + +# #ifdef __cplusplus +# extern "C" +# #endif +# void _mlir_gemm__mlir_ciface_cutlass_host_func_Tensorgenericoi641_Tensorgenericoi641_Tensorgenericoi641(void **args, int32_t num_args); + +# static inline void gemm_wrapper(gemm_Kernel_Metadata_t *metadata, gemm_Tensor_a_t *a, gemm_Tensor_b_t *b, gemm_Tensor_c_t *c) { +# void *args[4] = { +# a, b, c, &metadata->kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0 +# }; +# _mlir_gemm__mlir_ciface_cutlass_host_func_Tensorgenericoi641_Tensorgenericoi641_Tensorgenericoi641(args, 4); +# } diff --git a/python/CuTeDSL/cutlass/base_dsl/export/export.py b/python/CuTeDSL/cutlass/base_dsl/export/export.py new file mode 100644 index 00000000..836e658d --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/export/export.py @@ -0,0 +1,199 @@ +# 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 +# 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. + +import io +import os +import array +import tempfile + +from ..common import DSLRuntimeError +from ..jit_executor import JitCompiledFunction, get_escaped_cubin_bytes +from ...base_dsl.dsl import BaseDSL +from ...base_dsl.typing import Int32, Int64, Float32, Float64 +from ..._mlir import ir +from ..._mlir.dialects import llvm + +from .c_header_generator import CHeaderGenerator + +from typing import Union + +cubin_suffix = "cubin" + +def get_export_module(ir_module: ir.Module, symbol_prefix: str): + """Get the export module which is cloned from the original compiled ir module, and add the prefix + to avoid the symbol conflict. + + @param ir_module: The original compiled ir module. Comes from the JitCompiledFunction.ir_module. + @param symbol_prefix: The prefix name of the function. This is the unique identifier name of the function to avoid symbol conflict in the generated object file. + @return: The export module of the function. + """ + # Add prefix for symbol names to avoid conflict with other functions + defined_symbols = set() + + def walk_llvm_func_op(op): + # not a declaration + if ( + op.name == "llvm.func" + and len(op.opview.operation.regions) > 0 + and len(op.opview.operation.regions[0].blocks) > 0 + ): + defined_symbols.add(op.attributes["sym_name"].value) + op.attributes["sym_name"] = ir.StringAttr.get( + symbol_prefix + "_" + op.attributes["sym_name"].value + ) + return ir.WalkResult.ADVANCE + + def walk_llvm_call_op(op): + if op.name == "llvm.call" and op.attributes["callee"].value in defined_symbols: + op.attributes["callee"] = ir.FlatSymbolRefAttr.get( + symbol_prefix + "_" + op.attributes["callee"].value + ) + return ir.WalkResult.ADVANCE + + with ir.Context(): + export_module = ir.Module.parse(str(ir_module)) + export_module.operation.walk(walk_llvm_func_op) + export_module.operation.walk(walk_llvm_call_op) + return export_module + + + +def dump_to_object( + prefix_name: str, + export_module: ir.Module, + jit_function: Union[JitCompiledFunction, "CudaDialectJitCompiledFunction"], + dsl: BaseDSL, + use_gpu_dialect: bool, +) -> bytes: + """Dump the compiled ir function to a bytes object with ELF format. The bytes object contains the host + launch entry function and cubin inside. + + @param prefix_name: The prefix name of the function. This is the unique identifier name of the function to avoid symbol conflict in the generated object file. + @param export_module: The export module of the function. This is the module that contains the function to be exported. + @param jit_function: The jit-compiled function. To provided other metadata for the object file. + @param dsl: The dsl object. This is the dsl object to get the compiler provider and shared libs. + @return: The bytes object of the function. + """ + if use_gpu_dialect: + cubin_data = None + + def strip_gpu_binary_op(op): + if op.name == "gpu.binary": + s = io.BytesIO() + op.operation.write_bytecode(s) + nonlocal cubin_data + cubin_data = s.getvalue() + cubin_data = cubin_data.split(b'bin = "')[1].split(b'">')[0] + cubin_data = get_escaped_cubin_bytes(cubin_data) + op.erase() + return ir.WalkResult.ADVANCE + return ir.WalkResult.ADVANCE + + # Strip gpu related to avoid the object file generating builtin module load/unload functions + export_module.operation.walk(strip_gpu_binary_op) + + cubin_array = array.array("b", cubin_data) + with ( + export_module.context, + ir.Location.unknown(), + ir.InsertionPoint(export_module.body), + ): + new_binary_global_op = llvm.GlobalOp( + sym_name="_".join([prefix_name, cubin_suffix]), + global_type=ir.Type.parse(f"!llvm.array<{len(cubin_array)} x i8>"), + linkage=ir.Attribute.parse("#llvm.linkage"), + value=ir.DenseIntElementsAttr.get(cubin_array), + constant=True, + ) + if "gpu.container_module" in export_module.operation.attributes: + del export_module.operation.attributes["gpu.container_module"] + # Generate the object file + export_engine = dsl.compiler_provider.jit( + export_module, shared_libs=dsl.get_shared_libs() + ) + # This lookup is necessary to make sure the compilation is done. + entry_func = export_engine.raw_lookup( + "_".join([prefix_name, jit_function.function_name]) + ) + if not entry_func: + raise DSLRuntimeError( + f"Execution engine cannot find the entry function {prefix_name}_{jit_function.function_name}" + ) + try: + with tempfile.NamedTemporaryFile() as tmp_object_file: + export_engine.dump_to_object_file(tmp_object_file.name) + with open(tmp_object_file.name, "rb") as f: + ret = f.read() + return ret + except Exception as e: + raise DSLRuntimeError(f"Error dumping object file: {e}") from e + + +def export_to_c( + jit_function: Union[JitCompiledFunction, "CudaDialectJitCompiledFunction"], + file_path: str, + file_name: str, + dsl: BaseDSL, + c_header_generator: CHeaderGenerator, + use_gpu_dialect: bool, +): + """Exports the jit-compiled function to a C compatible files(header/library). + This is used for c/cpp AOT support. + The `file_name` will be used as the symbol prefix of the generated functions, it is guaranteed by + the caller that the generated functions are unique. And the function will always overwrite the existing file. + + + The c header file is generated with following components: + 1. The host launch entry function. And the structure definitions of the arguments. + 2. The device metadata load/unload functions. + 3. The cubin data array and len. + + The library contains the binary of the underlying host launch entry function. + + @param jit_function: The jit-compiled function from `cute.compile`. + @param file_path: The path to the directory where the header and object files will be saved. + @param file_name: The name of the function. This is the unique identifier name of the function to avoid symbol conflict in the generated object file. + @param dsl: The dsl object. This is the dsl object to get the compiler provider and shared libs. + @param c_header_generator: The c header generator. This is the c header generator to generate the c header file. + """ + export_module = get_export_module(jit_function.ir_module, file_name) + # Generate the c header file + header_file_content = c_header_generator( + file_name, + export_module, + jit_function.args_spec, + jit_function.function_name, + jit_function.kernel_info, + jit_function.dynamic_args, + jit_function.dynamic_kwargs, + dsl.name, + ) + try: + with open(os.path.join(file_path, file_name + ".h"), "w") as f: + f.write(header_file_content) + except Exception as e: + raise DSLRuntimeError(f"Error writing header file: {e}") from e + + # Generate the object file + object_file_content = dump_to_object( + file_name, + export_module, + jit_function, + dsl, + use_gpu_dialect, + ) + try: + with open(os.path.join(file_path, file_name + ".o"), "wb") as f: + f.write(object_file_content) + except Exception as e: + raise DSLRuntimeError(f"Error writing object file: {e}") from e + + diff --git a/python/CuTeDSL/cutlass/base_dsl/jit_executor.py b/python/CuTeDSL/cutlass/base_dsl/jit_executor.py index bd4bf34b..0982953f 100644 --- a/python/CuTeDSL/cutlass/base_dsl/jit_executor.py +++ b/python/CuTeDSL/cutlass/base_dsl/jit_executor.py @@ -28,14 +28,13 @@ from .._mlir import ir # Local modules imports from . import typing as t -from .common import DSLRuntimeError +from .common import DSLRuntimeError, DSLCudaRuntimeError from .runtime import cuda as cuda_helpers from .runtime.jit_arg_adapters import JitArgAdapterRegistry, is_arg_spec_constexpr from .typing import get_c_pointers from .utils.logger import log from .utils.timer import timer - class CudaModuleAndKernel: """A loaded CUDA kernel and its metadata.""" @@ -143,13 +142,11 @@ class ExecutionArgs: self.args_spec = self.filter_runtime_arg_spec(spec) self.original_args_spec = spec - def generate_execution_args(self, args, kwargs): + def get_rectified_args(self, args, kwargs): """ - This function is the prune version of `generate_mlir_function_types` which only generates execution args - to get rid of mlir context. + This function is used to rectify the args and kwargs to a final runtime argument list according to the args_spec. """ args_spec = self.args_spec - # Process positional arguments with defaults rectified_args = list(args) if args_spec.defaults and len(args) < len(args_spec.args): @@ -163,7 +160,9 @@ class ExecutionArgs: rectified_args.append(v) # Process keyword arguments - rectified_kwargs = {k: v for k, v in kwargs.items() if k not in args_spec.args} + rectified_kwargs = collections.OrderedDict( + (k, v) for k, v in kwargs.items() if k not in args_spec.args + ) if args_spec.kwonlydefaults and len(rectified_kwargs) < len( args_spec.kwonlyargs ): @@ -182,10 +181,18 @@ class ExecutionArgs: "function signature kwonlyargs length": len(args_spec.kwonlyargs), }, ) + return rectified_args + list(rectified_kwargs.values()) + + def generate_execution_args(self, args, kwargs): + """ + This function is the prune version of `generate_mlir_function_types` which only generates execution args + to get rid of mlir context. + """ + args_spec = self.args_spec exe_args = [] adapted_args = [] - input_args = rectified_args + list(rectified_kwargs.values()) + input_args = self.get_rectified_args(args, kwargs) input_arg_names = args_spec.args + args_spec.kwonlyargs for arg, arg_name in zip(input_args, input_arg_names): # short-cut for args already converted @@ -243,7 +250,7 @@ class ExecutionArgs: runtime_kwonlydefaults = {} if arg_spec.kwonlyargs: - for kwarg in arg_spec.kwonlyargs: + for i, kwarg in enumerate(arg_spec.kwonlyargs): arg_type = arg_spec.annotations.get(kwarg, None) # Apply same filtering logic @@ -366,6 +373,8 @@ class JitModule: for m in set([m.cuda_module for m in self.cuda_modules]): cuda_helpers.unload_library(m) self.cuda_modules.clear() + except Exception as e: + pass finally: self._unloaded = True @@ -382,8 +391,8 @@ class JitExecutor: def __init__( self, - jit_module: JitModule, - exec_context: JitExecuteContext, + jit_module: Union[JitModule, "CudaDialectJitModule"], + exec_context: Optional[JitExecuteContext], jit_time_profiling: bool, ): # JitExecutor will keep JitCompiledFunction alive so that the underlying @@ -393,12 +402,25 @@ class JitExecutor: self.exec_context = exec_context self.profiler = timer(enable=jit_time_profiling) + # Get the cuda result type from the capi function. + # This is only set to i32 if CudaDialectJitModule is used. + cuda_result_type = self.jit_module.capi_func.restype + self.cuda_result = cuda_result_type() if cuda_result_type is not None else None + # Assume each execution args has type `c_void_p` to reduce the overhead of `ctypes.cast`. def _get_invoke_packed_args(self, exe_args): - exe_args += self.exec_context.kernel_functions_ptrs + # If expecting a cuda result, add a pointer to exe_args + if self.cuda_result is not None: + exe_args.append(ctypes.addressof(self.cuda_result)) + if self.exec_context is not None: + exe_args += self.exec_context.kernel_functions_ptrs packed_args = (ctypes.c_void_p * len(exe_args))() for argNum in range(len(exe_args)): - packed_args[argNum] = exe_args[argNum] + arg = exe_args[argNum] + if isinstance(arg, ctypes.c_void_p): + packed_args[argNum] = arg + else: + packed_args[argNum] = ctypes.c_void_p(arg).value return packed_args def generate_execution_args(self, *args, **kwargs): @@ -408,6 +430,17 @@ class JitExecutor: try: packed_args = self.profiler(self._get_invoke_packed_args)(exe_args) self.profiler(self.jit_module.capi_func)(packed_args) + if self.cuda_result is not None: + if self.cuda_result.value != 0: + error_code = self.cuda_result.value + error_name = cuda_helpers._cudaGetErrorEnum( + cuda_helpers.cuda.CUresult(error_code) + ) + raise DSLCudaRuntimeError(error_code, error_name) + return self.cuda_result.value + return None + except DSLCudaRuntimeError as e: + raise e except Exception as e: raise DSLRuntimeError(f"💥💥💥 Runtime Crash 💥💥💥", cause=e) @@ -458,6 +491,8 @@ class JitCompiledFunction: kernel_info, jit_time_profiling, jit_function_artifacts, + prefix=None, + load_from_binary=False, ): self.ir_module = ir_module self.engine = engine @@ -473,6 +508,8 @@ class JitCompiledFunction: or jit_function_artifacts is None ) self.artifacts = jit_function_artifacts + self.prefix = prefix + self.load_from_binary = load_from_binary # This runtime state is stored here so that we can preserve the module # in the compiler cache. Callers can extend the lifetime of the module @@ -496,6 +533,35 @@ class JitCompiledFunction: """Returns the MLIR code of the JIT-compiled function.""" return self.artifacts.MLIR if self.artifacts is not None else None + def _deserializer(self): + """Load the cuda module from the binary execution engine. This function will be injected as the + JitCompiledFunction method which will be called by the jit executor to load the cuda module by AOT flow. + @param self: The JitCompiledFunction object. This is the JitCompiledFunction object to load the cuda module. + @param name: The name of the function. This is the unique identifier name of the function to avoid symbol conflict in the generated object file. + @param execution_engine: The binary execution engine. This is the execution engine to load the cuda module. + @param kernel_info: The kernel info. This is the kernel info to load the cuda module. + @return: The list of cuda modules. + """ + cubin_suffix = "cubin" + if self.prefix is None: + raise DSLRuntimeError("prefix is required to be set for binary loading") + cubin_data = self.engine.lookup("_".join([self.prefix, cubin_suffix])) + if not cubin_data: + raise RuntimeError( + "Unknown function " + "_".join([self.prefix, cubin_suffix]) + ) + cubin_module = cuda_helpers.load_library_data(cubin_data) + # load cuda module/get function pointer from module and cache + kernel_modules = collections.OrderedDict() + for sym, attrs in self.kernel_info.items(): + kernel = cuda_helpers.get_library_kernel(cubin_module, sym) + if cuda_helpers.get_driver_version() >= 11080: + attrs[ + cuda_helpers.cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED + ] = 1 + kernel_modules[sym] = CudaModuleAndKernel(sym, cubin_module, kernel, attrs) + return list(kernel_modules.values()) + def to(self, device=None) -> JitExecutor: """Returns an executable function bound to the given device. @@ -510,9 +576,10 @@ class JitCompiledFunction: with self._executor_lock: # We need to ensure that the modules are loaded if not already if self.jit_module is None: - cuda_modules = load_kernels_from_ir_module( - self.ir_module, self.kernel_info - ) + if self.ir_module is not None: + cuda_modules = load_kernels_from_ir_module( + self.ir_module, self.kernel_info + ) self.jit_module = JitModule( self.engine, self.capi_func, self.args_spec, cuda_modules ) @@ -522,6 +589,11 @@ class JitCompiledFunction: context = self.jit_module.get_device_execute_context(device) return JitExecutor(self.jit_module, context, self.jit_time_profiling) + def set_dynamic_args(self, dynamic_args, dynamic_kwargs): + """Sets the dynamic argument information required for export to c code generation.""" + self.dynamic_args = dynamic_args + self.dynamic_kwargs = dynamic_kwargs + def generate_execution_args(self, *args, **kwargs): return self.args_spec.generate_execution_args(args, kwargs) diff --git a/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py b/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py index 364333dd..64bfa2f9 100644 --- a/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py +++ b/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py @@ -23,15 +23,9 @@ import ctypes import cuda.bindings.driver as cuda import cuda.bindings.nvrtc as nvrtc -# MLIR imports -from ..._mlir import ir -from ..._mlir.dialects import gpu - # Local module imports from ..utils.logger import log as _log from ..common import * -from .jit_arg_adapters import JitArgAdapterRegistry - # ============================================================================= # Utils @@ -504,18 +498,18 @@ def load_library_data(cubin_data): """ Loads a CUBIN from data and returns the library. :param cubin_data: The binary data of the CUBIN. - :type cubin_data: bytes + :type cubin_data: bytes or ctypes.c_void_p :return: The library. :rtype: cuda.CUlibrary :raise DSLRuntimeError: If the CUDA operation fails. """ # Load module data - _log().info(f"cuLibraryLoadData {np.char.array(cubin_data).ctypes.data}") + if isinstance(cubin_data, bytes): + cubin_data = np.char.array(cubin_data).ctypes.data + _log().info(f"cuLibraryLoadData {cubin_data}") library = checkCudaErrors( - cuda.cuLibraryLoadData( - np.char.array(cubin_data).ctypes.data, None, None, 0, None, None, 0 - ) + cuda.cuLibraryLoadData(cubin_data, None, None, 0, None, None, 0) ) return library @@ -803,24 +797,3 @@ def get_device_attribute(attribute, device_id: int = 0): """ device = checkCudaErrors(cuda.cuDeviceGet(device_id)) return checkCudaErrors(cuda.cuDeviceGetAttribute(attribute, device)) - - -@JitArgAdapterRegistry.register_jit_arg_adapter(cuda.CUstream) -class StreamAdapter: - """ - Convert a CUDA stream to a stream representation for JIT arg generation. - """ - - def __init__(self, arg): - self._arg = arg - self._c_pointer = self._arg.getPtr() - - def __new_from_mlir_values__(self, values): - assert len(values) == 1 - return values[0] - - def __c_pointers__(self): - return [self._c_pointer] - - def __get_mlir_types__(self): - return [gpu.AsyncTokenType.get()] diff --git a/python/CuTeDSL/cutlass/base_dsl/runtime/stream_adapter.py b/python/CuTeDSL/cutlass/base_dsl/runtime/stream_adapter.py new file mode 100644 index 00000000..d302283c --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/runtime/stream_adapter.py @@ -0,0 +1,43 @@ +# 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 +# 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. + +""" +This module provides CUDA Python helper functions +""" + +import cuda.bindings.driver as cuda + +# MLIR imports +from ..._mlir import ir +from ..._mlir.dialects import gpu + +from .jit_arg_adapters import JitArgAdapterRegistry + + +@JitArgAdapterRegistry.register_jit_arg_adapter(cuda.CUstream) +class StreamAdapter: + """ + Convert a CUDA stream to a stream representation for JIT arg generation. + """ + + def __init__(self, arg): + self._arg = arg + self._c_pointer = self._arg.getPtr() + + def __new_from_mlir_values__(self, values): + assert len(values) == 1 + return values[0] + + def __c_pointers__(self): + return [self._c_pointer] + + def __get_mlir_types__(self): + return [gpu.AsyncTokenType.get()] diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/README.md b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/README.md new file mode 100644 index 00000000..9104592d --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/README.md @@ -0,0 +1,84 @@ +# TVM FFI DSL Bridge + +This folder implements a binding generator that +can be used by DSL compilers to generate functions in tvm-ffi ABI. + +- Interface specification: we provide a spec namespace that allows + developers to define the interface of the function, with possibly shape and type constraints. +- Constraint checking: we will generate MLIR functions that exposes the tvm-ffi ABI convention + checks the overall constraints and reads all related variable fields. +- User defined calling convention: user can provide a `CallProvider` instance + that takes `CallContext` and emit a call to internal DSL function. + +## Basic Example: Mixed Parameters + +```python + +from cutlass.base_dsl.tvm_ffi_builder import ( + spec, attach_ffi_func, ExecutionEngine, NopProvider +) +from cutlass._mlir import ir +import tvm_ffi +import numpy as np + +# Define parameters: int + tensor with symbolic shape +n = spec.Var("n", "int32") +params = [ + spec.Var("batch_size", "int32"), # Integer parameter + spec.Tensor("data", [n, 128], "float32"), # Tensor with symbolic first dimension +] + +# Generate function +# Function signature: process_data(batch_size: int32, data: Tensor([n, 128], float32)) +with ir.Context(), ir.Location.unknown(): + module = ir.Module.create() + attach_ffi_func(module, "process_data", params, NopProvider()) + + # Execute + engine = ExecutionEngine(module, opt_level=2, shared_libs=[]) + func = tvm_ffi.Function.__from_mlir_packed_safe_call__( + engine.raw_lookup("__tvm_ffi_process_data") + ) + + # Call function + data = tvm_ffi.from_dlpack(np.zeros((10, 128), dtype=np.float32)) + func(42, data) # batch_size=42, data shape=[10, 128] +``` + +## Matrix Multiplication Example + +```python +# Define matrix multiplication with automatic constraint validation +n = spec.Var("n", "int32") +m = spec.Var("m", "int32") +k = spec.Var("k", "int32") + +with spec.DefaultConfig(device="cpu"): + params = [ + spec.Tensor("A", [n, k], "float32"), # A: n×k + spec.Tensor("B", [k, m], "float32"), # B: k×m + spec.Tensor("C", [n, m], "float32"), # C: n×m + ] + +# Generate function +# Function signature: matmul(A: Tensor([n, k], float32), B: Tensor([k, m], float32), C: Tensor([n, m], float32)) +with ir.Context(), ir.Location.unknown(): + module = ir.Module.create() + attach_ffi_func(module, "matmul", params, NopProvider()) + + # Execute + engine = ExecutionEngine(module, opt_level=2, shared_libs=[]) + func = tvm_ffi.Function.__from_mlir_packed_safe_call__( + engine.raw_lookup("__tvm_ffi_matmul") + ) + + # Valid call: 2×3 × 3×4 = 2×4 + A = tvm_ffi.from_dlpack(np.zeros((2, 3), dtype=np.float32)) + B = tvm_ffi.from_dlpack(np.zeros((3, 4), dtype=np.float32)) + C = tvm_ffi.from_dlpack(np.zeros((2, 4), dtype=np.float32)) + func(A, B, C) # ✅ Valid dimensions + + # Invalid call: dimension mismatch + A_wrong = tvm_ffi.from_dlpack(np.zeros((2, 4), dtype=np.float32)) # Wrong: 2×4 + func(A_wrong, B, C) # ❌ Error: A[1]=4 != B[0]=3 +``` diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/__init__.py b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/__init__.py new file mode 100644 index 00000000..d0608a2f --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/__init__.py @@ -0,0 +1,32 @@ +# 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 +# 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. + +"""Helper tool to build TVM-FFI functions using MLIR.""" + +from .call_provider import DynamicParamPackCallProvider, NopCallProvider +from .spec import Param, Var +from .tvm_ffi_builder import ( + CallContext, + CallProvider, + attach_ffi_func, + rename_tvm_ffi_function, +) + +__all__ = [ + "CallContext", + "CallProvider", + "DynamicParamPackCallProvider", + "NopCallProvider", + "Param", + "Var", + "attach_ffi_func", + "rename_tvm_ffi_function", +] diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/call_provider.py b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/call_provider.py new file mode 100644 index 00000000..03512528 --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/call_provider.py @@ -0,0 +1,245 @@ +# 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 +# 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. + +"""Call provider that implements a specific calling convention.""" + +from dataclasses import dataclass +from typing import Optional, Union + +from . import spec +from ..._mlir import ir +from ..._mlir.dialects import llvm +from .tvm_ffi_builder import CallContext, CallProvider, TVMFFIBuilder + + +class NopCallProvider(CallProvider): + """No-op call provider for testing purposes.""" + + def __call__(self, current_block: ir.Block, context: CallContext) -> ir.Block: + """No-op call provider that just returns the current block.""" + return current_block + + +class DynamicParamPackCallProvider(CallProvider, TVMFFIBuilder): + """Packs dynamic arguments to a struct then calls the function. + + .. code-block:: c + + void call(Tensor0 t0, Tensor1 t1) { + // packed arguments + void** packed_args[] = {&t0, &t1}; + // call target + target_func(packed_args); + } + + Parameters + ---------- + target_func: str + The name of the target function. + + include_num_args: bool + Whether to include the number of arguments in the packed arguments. + + struct_call: bool + Whether to use the struct call convention. + """ + + def __init__( + self, target_func: str, include_num_args: bool = False, struct_call: bool = False, + ) -> None: + import tvm_ffi + + TVMFFIBuilder.__init__(self) + self.target_func = target_func + self.include_num_args = include_num_args + self.struct_call = struct_call + self.float4x2_dtype = tvm_ffi.dtype("float4_e2m1fnx2") + + def get_callee_struct_for_param_tensor( + self, + param: spec.Tensor, + current_block: ir.Block, + data: ir.Value, + shape: list[ir.Value], + strides: list[ir.Value], + flatten_struct: ir.Type, + ) -> ir.Type: + """Routine used to override tensor passsing struct conention.""" + return flatten_struct + + def pack_param_tensor( + self, current_block: ir.Block, context: CallContext, param: spec.Tensor + ) -> tuple[ir.Type, ir.Value]: + """Pack a tensor parameter to a struct.""" + map_shape_value = lambda _, value: value + map_stride_value = lambda _, value: value + + if param.map_tensor_dtype_f4x2_to_f4 and param.dtype == self.float4x2_dtype: + # specially handle f4x2 to f4 tensor conversion + # we multiply the stride by 2 for all dimensions except the one with stride=1 + # we also multiply the shape by 2 for the specific dimension with stride=1 + def find_stride_one_index() -> int: + if param.strides is None: + return len(param.shape) - 1 + for i, stride in enumerate(param.strides): + if isinstance(stride, int) and stride == 1: + return i + raise ValueError("stride=1 index not found, needed for f4 tensor") + stride_one_index = find_stride_one_index() + + def map_shape_for_tensor_dtype_f4x2_to_f4(index: int, value: ir.Value) -> ir.Value: + if index == stride_one_index: + with ir.InsertionPoint(current_block): + return self.mul(value, self.integer_constant(value.type, 2)) + return value + + def map_stride_for_tensor_dtype_f4x2_to_f4(index, value: ir.Value) -> ir.Value: + if index != stride_one_index: + with ir.InsertionPoint(current_block): + return self.mul(value, self.integer_constant(value.type, 2)) + return value + + map_shape_value = map_shape_for_tensor_dtype_f4x2_to_f4 + map_stride_value = map_stride_for_tensor_dtype_f4x2_to_f4 + + data = context.matched_var_binding[param.data] + shape = [] + strides = [] + # append all vars in shape + for index, dim in enumerate(param.shape): + if isinstance(dim, spec.Var): + shape.append(map_shape_value(index, context.matched_var_binding[dim])) + # append all vars in strides + if param.strides is not None: + for index, dim in enumerate(param.strides): + if isinstance(dim, spec.Var): + strides.append(map_stride_value(index, context.matched_var_binding[dim])) + flatten_struct, alloca = self.pack_values_to_alloca( + current_block, context.entry_block, [data, *shape, *strides] + ) + callee_struct = self.get_callee_struct_for_param_tensor( + param, current_block, data, shape, strides, flatten_struct + ) + + return callee_struct, alloca + + def pack_param_var( + self, current_block: ir.Block, context: CallContext, param: spec.Var + ) -> tuple[ir.Type, ir.Value]: + """Pack a var parameter to a struct.""" + value: ir.Value = context.matched_var_binding[param] + _, alloca = self.pack_values_to_alloca( + current_block, context.entry_block, [value] + ) + return (value.type, alloca) + + def pack_param_shape( + self, current_block: ir.Block, context: CallContext, param: spec.Shape + ) -> tuple[ir.Type, ir.Value]: + """Pack a shape parameter to a struct.""" + dynamic_args: list[ir.Value] = [] + for dim in param.shape: + if isinstance(dim, spec.Var): + dynamic_args.append(context.matched_var_binding[dim]) + return self.pack_values_to_alloca( + current_block, context.entry_block, dynamic_args + ) + + def pack_params( + self, current_block: ir.Block, context: CallContext + ) -> list[tuple[ir.Type, ir.Value]]: + """Pack a parameter to a struct.""" + packed_params = [] + for param in context.params: + if isinstance(param, spec.Tensor): + packed_params.append( + self.pack_param_tensor(current_block, context, param) + ) + elif isinstance(param, spec.Var): + packed_params.append(self.pack_param_var(current_block, context, param)) + elif isinstance(param, spec.Shape): + packed_params.append( + self.pack_param_shape(current_block, context, param) + ) + elif isinstance(param, (spec.Stream, spec.EnvStream)): + packed_params.append( + self.pack_param_var(current_block, context, param.var) + ) + elif isinstance(param, spec.DataPointer): + packed_params.append( + self.pack_param_var(current_block, context, param.var) + ) + else: + raise NotImplementedError(f"Unsupported parameter type: {type(param)}") + return packed_params + + def generate_llvm_call( + self, + current_block: ir.Block, + call_operands: list[ir.Value], + context: CallContext, + ) -> ir.Block: + """Generate the LLVM call operation.""" + with ir.InsertionPoint(current_block): + llvm.call( + result=None, + callee=self.target_func, + callee_operands=call_operands, + op_bundle_sizes=[], + op_bundle_operands=[], + ) + return current_block + + def load_to_call_operands( + self, + struct_type: Union[ir.Type, tuple[ir.Type]], + alloca: Union[ir.Value, tuple[ir.Value]], + ) -> list[ir.Value]: + """Load the packed parameters to the call operands.""" + assert (isinstance(struct_type, ir.Type) and isinstance(alloca, ir.Value)) or ( + isinstance(struct_type, tuple) and isinstance(alloca, tuple) + ) + if isinstance(struct_type, tuple): + return [ + llvm.load(struct_type[i], alloca[i]) for i in range(len(struct_type)) + ] + return [llvm.load(struct_type, alloca)] + + def __call__(self, current_block: ir.Block, context: CallContext) -> ir.Block: + """Alloca call provider that uses dynamic param pack call convention.""" + packed_params = self.pack_params(current_block, context) + + if self.struct_call: + # load back arguments as structs from alloca + call_operands = [] + with ir.InsertionPoint(current_block): + for struct_type, alloca in packed_params: + call_operands += self.load_to_call_operands(struct_type, alloca) + else: + # pack the values to an alloca that we can pass as void** + all_values = [] + for _, value in packed_params: + if isinstance(value, tuple): + all_values.extend(value) + else: + all_values.append(value) + _, packed_args_value = self.pack_values_to_alloca( + current_block, context.entry_block, all_values + ) + + call_operands = [packed_args_value] + if self.include_num_args: + with ir.InsertionPoint(current_block): + num_args = self.i32(len(all_values)) + call_operands.append(num_args) + + current_block = self.generate_llvm_call(current_block, call_operands, context) + return current_block 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 new file mode 100644 index 00000000..943ce262 --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py @@ -0,0 +1,468 @@ +# 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 +# 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. + +"""MLIR type builder and basic operations for LLVM dialect.""" + +from collections.abc import Sequence +from typing import Any, Optional + +from ..._mlir import ir +from ..._mlir.dialects import llvm + + +class MLIRTypeBuilder: + """Builder for MLIR types and basic operations.""" + + BRANCH_WEIGHTS_LIKELY = (2000, 1) + BRANCH_WEIGHTS_UNLIKELY = (1, 2000) + + def __init__(self) -> None: + """Initialize the MLIR type builder.""" + self.i32_type = ir.IntegerType.get_signless(32) + self.ui32_type = ir.IntegerType.get_unsigned(32) + self.i64_type = ir.IntegerType.get_signless(64) + self.i16_type = ir.IntegerType.get_signless(16) + self.i8_type = ir.IntegerType.get_signless(8) + self.i1_type = ir.IntegerType.get_signless(1) + self.f32_type = ir.Type.parse("f32") + self.f64_type = ir.Type.parse("f64") + self.ptr_type = llvm.PointerType.get() + self.gpu_ptr_type = llvm.PointerType.get(address_space=1) + # did not find a programmatic way to get the void type + self.void_type = ir.Type.parse("!llvm.void") + + def ptr_type_with_address_space( + self, address_space: Optional[int] = None + ) -> ir.Type: + """Get the pointer type with the given address space.""" + if address_space is None or address_space == 0: + return self.ptr_type + return llvm.PointerType.get(address_space=address_space) + + def as_attr(self, tp: ir.Type) -> ir.TypeAttr: + """Convert the type to a type attribute.""" + return ir.TypeAttr.get(tp) + + def int_type(self, bits: int) -> ir.Type: + """Get the `i` type.""" + return ir.IntegerType.get_signless(bits) + + def uint_type(self, bits: int) -> ir.Type: + """Get the `ui` type.""" + return ir.IntegerType.get_unsigned(bits) + + def struct_type( + self, + *, + name: Optional[str] = None, + fields: Sequence[ir.Type] = (), + packed: bool = False, + ) -> ir.Type: + """Get or create a struct type. + + Parameters + ---------- + name : Optional[str] + The name of the struct type. If not provided, a `literal` struct type is created, + which is identified by its fields only. + fields : Sequence[ir.Type] + The fields of the struct type. + packed : bool + Whether to create a packed struct type. If `True`, there is no padding between fields. + Otherwise, there might be padding + between fields to ensure alignment. + + See Also + -------- + https://mlir.llvm.org/docs/Dialects/LLVM/#structure-types + + """ + if name is None: + return llvm.StructType.get_literal(fields, packed=packed) + else: + return llvm.StructType.new_identified(name, fields, packed=packed) + + def identified_struct_type(self, name: str) -> ir.Type: + """Get a previously created identified struct type by its name.""" + return llvm.StructType.get_identified(name) + + def func_type(self, *, params: Sequence[ir.Type] = (), ret: ir.Type) -> ir.Type: + """Get a function type. + + Parameters + ---------- + params : Sequence[ir.Type] + The parameters of the function type. + ret : ir.Type + The return type of the function type. + + See Also + -------- + https://mlir.llvm.org/docs/Dialects/LLVM/#function-types + + """ + return ir.Type.parse( + "!llvm.func<{} ({})>".format(str(ret), ", ".join(map(str, params))) + ) # did not find a programmatic way to get the function type + + def global_dtor_entry_type(self) -> ir.Type: + """Get the type of the global destructor entry. + + See Also: + - https://llvm.org/docs/LangRef.html#the-llvm-global-dtors-global-variable + - https://mlir.llvm.org/docs/Dialects/LLVM/#llvmmlirglobal_dtors-llvmglobaldtorsop + """ + return self.struct_type(fields=[self.i32_type, self.ptr_type, self.ptr_type]) + + +class MLIRBuilder(MLIRTypeBuilder): + """A builder for MLIR related types and operations. + + Convention: + all statement-generation methods expect we are inside a insertion point of the current_block. + If the statement terminates the current block, it will assign the new block to the current_block + but not set the insersion point. + """ + + def __init__(self) -> None: + """Initialize the MLIR builder.""" + super().__init__() + self.module: Optional[ir.Module] = None + self.const_str_table: dict[str, ir.Value] = {} + self.get_element_extra_kwargs: dict[str, Any] = {} + + # create constants + def integer_constant(self, tp: ir.Type, value: int) -> ir.Value: + """Create an integer constant with the given type and value.""" + return llvm.ConstantOp(tp, ir.IntegerAttr.get(tp, value)).res + + def i32(self, value: int) -> ir.Value: + """Create an i32 constant with the given value.""" + return self.integer_constant(self.i32_type, value) + + def ui32(self, value: int) -> ir.Value: + """Create a ui32 constant with the given value.""" + return self.integer_constant(self.ui32_type, value) + + def i1(self, value: int) -> ir.Value: + """Create an i1 constant with the given value.""" + return self.integer_constant(self.i1_type, value) + + def i8(self, value: int) -> ir.Value: + """Create an i8 constant with the given value.""" + return self.integer_constant(self.i8_type, value) + + def i16(self, value: int) -> ir.Value: + """Create an i16 constant with the given value.""" + return self.integer_constant(self.i16_type, value) + + def i64(self, value: int) -> ir.Value: + """Create an i64 constant with the given value.""" + return self.integer_constant(self.i64_type, value) + + def mul(self, lhs: ir.Value, rhs: ir.Value) -> ir.Value: + """Create a multiplication operation between two values.""" + return llvm.mul(lhs, rhs, overflow_flags=llvm.IntegerOverflowFlags.none) + + # expressions + def not_equal(self, lhs: ir.Value, rhs: ir.Value) -> ir.Value: + """Create a not-equal comparison between two values.""" + return llvm.icmp(llvm.ICmpPredicate.ne, lhs, rhs) + + def equal(self, lhs: ir.Value, rhs: ir.Value) -> ir.Value: + """Create an equal comparison between two values.""" + return llvm.icmp(llvm.ICmpPredicate.eq, lhs, rhs) + + def or_(self, lhs: ir.Value, rhs: ir.Value) -> ir.Value: + """Create a logical OR operation between two values.""" + return llvm.or_(lhs, rhs) + + def and_(self, lhs: ir.Value, rhs: ir.Value) -> ir.Value: + """Create a logical AND operation between two values.""" + return llvm.and_(lhs, rhs) + + def not_(self, value: ir.Value) -> ir.Value: + """Create a logical NOT operation.""" + # Ensure we're working with i1 type for boolean operations + if value.type != self.i1_type: + value = llvm.trunc(res=self.i1_type, arg=value) + return llvm.xor(value, self.i1(1)) + + def i64_divisible_const(self, value: ir.Value, align_const: int) -> ir.Value: + """Check if i64 value is divisible by align_const. + + Parameters + ---------- + value : ir.Value + The i64 value to check. + align_const : int + The alignment constant to check divisibility against. + + Returns + ------- + ir.Value + A boolean value (i1) indicating if value is divisible by align_const. + + Notes + ----- + Uses fast path (bitwise AND) when align_const is a power of two, + otherwise uses modulo operation. + """ + # Check if align_const is a power of two + is_power_of_two = (align_const > 0) and (align_const & (align_const - 1)) == 0 + + if is_power_of_two: + # Fast path: use bitwise AND + # value is divisible by align_const iff (value & (align_const - 1)) == 0 + mask = self.i64(align_const - 1) + masked = llvm.and_(value, mask) + return self.equal(masked, self.i64(0)) + else: + # Slow path: use modulo operation + align_val = self.i64(align_const) + remainder = llvm.urem(value, align_val) + return self.equal(remainder, self.i64(0)) + + def br(self, target_block: ir.Block, *, args: Optional[list[ir.Value]] = None) -> None: + """Create an unconditional branch. + + Parameters + ---------- + target_block : ir.Block + The target block to branch to. + args : list[ir.Value], optional + The values to pass as arguments to the target block. If None, + no arguments are passed. The target block must have the same + number of arguments as values in this list. + + """ + if args is None: + llvm.br(dest=target_block, dest_operands=[]) + else: + llvm.br(dest_operands=args, dest=target_block) + + def address_of(self, name: str, tp: ir.Type) -> ir.Value: + """Get the address of a global symbol.""" + return llvm.AddressOfOp(tp, name).result + + def getelementptr( + self, ptr: ir.Value, constant_indices: Sequence[int], elem_type: ir.Type + ) -> ir.Value: + """Create a getelementptr operation. + + Parameters + ---------- + ptr : ir.Value + The pointer to the element. + indices : Sequence[ir.Value] + The indices to the element. + elem_type : ir.Type + The type of the element. + + Returns + ------- + ir.Value + The resulting element pointer. + """ + try: + return llvm.getelementptr( + self.ptr_type, + ptr, + [], + raw_constant_indices=ir.DenseI32ArrayAttr.get(constant_indices), + elem_type=elem_type, + **self.get_element_extra_kwargs, + ) + except Exception: + # compatibility with different LLVM versions + self.get_element_extra_kwargs = {"no_wrap_flags": []} + return llvm.getelementptr( + self.ptr_type, + ptr, + [], + raw_constant_indices=ir.DenseI32ArrayAttr.get(constant_indices), + elem_type=elem_type, + **self.get_element_extra_kwargs, + ) + + def return_(self, ret: Optional[ir.Value] = None) -> None: + """Create a return statement.""" + llvm.return_(arg=ret) + + def cond_br( + self, + cond: ir.Value, + true_block: ir.Block, + false_block: ir.Block, + branch_weights=None + ) -> None: + """Create a conditional branch. + + Parameters + ---------- + cond : ir.Value + The condition value (i1 type). + true_block : ir.Block + The block to branch to if condition is true. + false_block : ir.Block + The block to branch to if condition is false. + branch_weights : Optional[tuple[int, int]] + Optional branch weights [true_weight, false_weight] for optimization hints. + Higher values indicate higher probability. For example, (99, 1) indicates + the true branch is much more likely than the false branch. + """ + if branch_weights is not None: + # Branch weights should be a tuple/list of two integers [true_weight, false_weight] + if len(branch_weights) != 2: + raise ValueError("branch_weights must have exactly 2 elements") + llvm.cond_br( + cond, + true_dest_operands=[], + false_dest_operands=[], + true_dest=true_block, + false_dest=false_block, + branch_weights=ir.DenseI32ArrayAttr.get(list(branch_weights)) + ) + else: + llvm.cond_br( + cond, + true_dest_operands=[], + false_dest_operands=[], + true_dest=true_block, + false_dest=false_block, + ) + + def define_global_string(self, content: str) -> str: + """Define a global string symbol with the given content using standard MLIR APIs.""" + if content in self.const_str_table: + return self.const_str_table[content] + symbol = f"__tvm_ffi__str_{len(self.const_str_table)}" + + # The string_attr.value gives us the original string, but we need to escape it for + # MLIR parsing. Let's use a simple approach: escape quotes and backslashes + escaped_content = content.replace("\\", "\\\\").replace('"', '\\"') + + # Parse the MLIR string with proper escaping using the standard API + module_body = self.module.body # type: ignore[union-attr] + with ir.InsertionPoint(module_body): + parsed_op = ir.Operation.parse( + f'llvm.mlir.global private constant @{symbol}("{escaped_content}\\00")' + ) + module_body.append(parsed_op) + self.const_str_table[content] = symbol + return symbol + + # function + def function( + self, + name: str, + params_type: Sequence[ir.Type], + ret_type: ir.Type, + internal: bool = False, + ) -> tuple[list[ir.Value], ir.Block]: + """Create a function with the given signature.""" + func_op = llvm.func( + name, + function_type=self.as_attr( + self.func_type(ret=ret_type, params=params_type) + ), + ) + if internal: + func_op.attributes["llvm.linkage"] = ir.StringAttr.get("private") + else: + func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get() + + params = [] + func_body: Any = func_op.body + if func_body is not None: + entry_block = ir.Block.create_at_start(func_body) + else: + raise RuntimeError("Function body is None") + for param_type in params_type: + params.append(entry_block.add_argument(param_type, ir.Location.unknown())) + return params, entry_block + + def declare_extern_func( + self, name: str, params: Sequence[ir.Type], ret: ir.Type + ) -> None: + """Define extern function.""" + func_type = self.func_type(params=params, ret=ret) + func_op = llvm.func( + name, + function_type=self.as_attr(func_type), + ) + func_op.attributes["llvm.linkage"] = ir.StringAttr.get("external") + + def pack_values_to_alloca( + self, + current_block: ir.Block, + entry_block: ir.Block, + values: Sequence[ir.Value], + ) -> tuple[ir.Type, ir.Value]: + """Pack values to an alloca that lays out in the order of the values. + + Parameters + ---------- + current_block : ir.Block + The current block. + entry_block : ir.Block + The entry block to create an alloca for the struct. + values : Sequence[ir.Value] + The values to pack. + + Returns + ------- + tuple[ir.Type, ir.Value] + The struct type and the alloca. + """ + with ir.InsertionPoint(entry_block.operations[0]): + # declare the struct type + struct_type = self.struct_type(fields=[value.type for value in values]) + alloca = llvm.alloca( + res=self.ptr_type, + elem_type=struct_type, + array_size=self.i32(1), + ) + + with ir.InsertionPoint(current_block): + for index, value in enumerate(values): + # Get pointer to the field at the given index + field_ptr = self.getelementptr( + alloca, + [0, index], + struct_type, + ) + # Store the value into the field + llvm.store(value, field_ptr) + return (struct_type, alloca) + + def find_operations_in_module( + self, module: ir.Module, name: str + ) -> list[ir.Operation]: + """Find operations in the module by the operation name.""" + operations = [] + for op in module.body: # type: ignore[union-attr] + if op.name == name: + operations.append(op) + return operations + + def find_func_in_module( + self, module: ir.Module, name: str + ) -> Optional[ir.Operation]: + """Find a function in the module.""" + for op in module.body: # type: ignore[union-attr] + if op.name == "llvm.func": + # Get the function name from the sym_name attribute + if "sym_name" in op.attributes: + func_name = str(op.attributes["sym_name"]).strip('"') + if func_name == name: + return op + return None diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/spec.py b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/spec.py new file mode 100644 index 00000000..62f7698b --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/spec.py @@ -0,0 +1,431 @@ +# 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 +# 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. + +"""Kernel specification classes for TVM-FFI function parameters.""" +from abc import ABC + +from collections.abc import Sequence +from typing import Optional, Union + +try: + import tvm_ffi +except ModuleNotFoundError: + pass + +class DefaultConfig: + """Default configuration with context manager support.""" + + _current: Optional["DefaultConfig"] = None + _old_current: Optional["DefaultConfig"] = None + device_type: str + + def __init__(self, *, device_type: Optional[str] = None) -> None: + """Initialize a default configuration. + + Parameters + ---------- + device_type : str, optional + The device type (e.g., "cpu", "cuda", "metal"). If None, copies from current config. + """ + if device_type is None: + device_type = DefaultConfig.current().device_type # type: ignore[union-attr] + self.device_type = device_type + + def __enter__(self) -> "DefaultConfig": + """Enter the context manager.""" + self._old_current = DefaultConfig._current + DefaultConfig._current = self + return self + + def __exit__( + self, + exc_type: Optional[type], + exc_val: Optional[BaseException], + exc_tb: Optional[object], + ) -> None: + """Exit the context manager.""" + DefaultConfig._current = self._old_current + + @classmethod + def current(cls) -> Optional["DefaultConfig"]: + """Get the current default configuration. + + Returns + ------- + Optional[DefaultConfig] + The current default configuration. + """ + return cls._current + + @classmethod + def _set_init_default_config(cls) -> None: + """Set the initial default configuration.""" + current = cls.__new__(cls) + current.device_type = "cuda" + cls._current = current + + +# Initialize the default config to cuda +DefaultConfig._set_init_default_config() + + +class Param(ABC): + """Base class for all parameters.""" + + +class Var(Param): + """variables: pointer, integer, floating-point, boolean, etc. + + Parameters + ---------- + name : str + The parameter name. + dtype : str | tvm_ffi.dtype + The data type of the parameter. + divisibility: Optional[int] + The divisibility of the parameter, by default None. + + """ + + name: str + dtype: "tvm_ffi.dtype" + divisibility: Optional[int] + + def __init__( + self, + name: str, + dtype: Union[str, "tvm_ffi.dtype"], + *, + divisibility: Optional[int] = None, + ) -> None: + """Initialize a Var parameter. + + Parameters + ---------- + name : str + The parameter name. + dtype : str | tvm_ffi.dtype + The data type of the parameter. + + """ + self.name = name + self.dtype = tvm_ffi.dtype(dtype) + self.divisibility = divisibility + + +class Shape(Param): + """Shape parameter. + + Parameters + ---------- + name : str + The parameter name. + shape: list[int | Var] + The shape of the parameter. + + """ + + name: str + shape: list[Union[int, Var]] + + def __init__(self, name: str, shape: list[Union[int, Var]]) -> None: + """Initialize a Shape parameter. + + Parameters + ---------- + name : str + The parameter name. + shape : list[int | Var] + The shape of the parameter. + + """ + self.name = name + self.shape = shape + + +class Tensor(Param): + """Tensor parameter. + + Parameters + ---------- + name : str + The parameter name. + dtype : str | tvm_ffi.dtype + The data type of the parameter. + shape : Sequence[int | Var] + The shape of the parameter. + device_type : int + The device type of the parameter. + device_id : Var + The device id of the parameter. + strides : Optional[Sequence[Var]], optional + The strides of the parameter, by default None. + """ + + name: str + shape: list[Union[int, Var]] + dtype: "tvm_ffi.dtype" + strides: Optional[list[Var]] + dlpack_device_type: int + device_id: Var + map_tensor_dtype_f4x2_to_f4: bool + data_alignment: Optional[int] + + def __init__( + self, + name: str, + shape: Sequence[Union[int, Var]], + dtype: Union[str, "tvm_ffi.dtype"], + *, + device_type: Optional[str] = None, + strides: Optional[Sequence[Var]] = None, + map_tensor_dtype_f4x2_to_f4: bool = False, + data_alignment: Optional[int] = None, + ) -> None: + """Initialize a Tensor parameter. + + Parameters + ---------- + name : str + The parameter name. + device_type : str + The device type of the parameter. + shape : Sequence[int | Var] + The shape of the parameter. + dtype : str | tvm_ffi.dtype + The data type of the parameter. + strides : Optional[Sequence[Var]], optional + The strides of the parameter, by default None. + map_tensor_dtype_f4x2_to_f4: bool + Whether to map tensor dtype float4x2 to float4 for internal use. + data_alignment: Optional[int], optional + The data alignment of the parameter, by default None. + """ + self.name = name + self.data = Var(name + ".data", tvm_ffi.dtype("handle")) + self.shape: list[Union[int, Var]] = list(shape) + self.dtype = tvm_ffi.dtype(dtype) + self.strides: Optional[list[Var]] = list(strides) if strides is not None else None + self.data_alignment = data_alignment + + # Use default device type if none specified + if device_type is None: + device_type = DefaultConfig.current().device_type # type: ignore[union-attr] + + example_device = tvm_ffi.device(device_type, 0) + self.dlpack_device_type = example_device.dlpack_device_type() + self.device_type_name = example_device.type + self.device_id = Var(name + ".device_id", tvm_ffi.dtype("int32")) + self.map_tensor_dtype_f4x2_to_f4 = map_tensor_dtype_f4x2_to_f4 + + +class Stream(Param): + """Stream parameter.""" + + name: str + var: Var + + def __init__(self, name: str) -> None: + """Initialize a Stream parameter. + + Parameters + ---------- + name : str + The parameter name. + """ + self.name = name + self.var = Var(name, tvm_ffi.dtype("handle")) + + +class EnvStream(Param): + """EnvStream parameter. + + Note + ---- + This parameter spec indicates that we expect + to call the function with TVMFFIEnvGetStream to get the stream + and it is not part of the FFI function signature. + """ + + name: str + var: Var + + def __init__(self, name: str) -> None: + """Initialize a EnvStream parameter. + + Parameters + ---------- + name : str + The parameter name. + """ + self.name = name + self.var = Var(name, tvm_ffi.dtype("handle")) + + +class DataPointer(Param): + """Data pointer parameter. + + The main difference between DataPointer and Var with handle type + is that DataPointer can contain address space information. + We also allow passing in int as data pointer, so it can + conveniently support torch.Tensor.data_ptr() as input. + + Parameters + ---------- + name : str + The parameter name. + """ + + name: str + var: Var + address_space: Optional[int] + + def __init__(self, name: str, address_space: Optional[int] = None) -> None: + """Initialize a DataPointer parameter. + + Parameters + ---------- + name : str + The parameter name. + """ + self.name = name + self.var = Var(name, tvm_ffi.dtype("handle")) + self.address_space = address_space + + +def signature(name: str, params: list[Param]) -> str: + """Generate a function signature string from name and parameters. + + Parameters + ---------- + name : str + The function name. + params : list[Param] + List of parameter objects (Var or Tensor). + + Returns + ------- + str + The formatted function signature. + + Raises + ------ + ValueError + If an unknown parameter type is encountered. + + """ + param_strs = [] + + for param in params: + if isinstance(param, Var): + param_str = f"{param.name}: {param.dtype}" + elif isinstance(param, Tensor): + # Format tensor shape + shape_strs = [] + for dim in param.shape: + if isinstance(dim, Var): + shape_strs.append(dim.name) + else: + shape_strs.append(str(dim)) + shape_str = "[" + ", ".join(shape_strs) + "]" + param_str = f"{param.name}: Tensor({shape_str}, {param.dtype})" + elif isinstance(param, Shape): + # Format shape parameter + shape_strs = [] + for dim in param.shape: + if isinstance(dim, Var): + shape_strs.append(dim.name) + else: + shape_strs.append(str(dim)) + shape_str = "[" + ", ".join(shape_strs) + "]" + param_str = f"{param.name}: Shape({shape_str})" + elif isinstance(param, Stream): + param_str = f"{param.name}: Stream" + elif isinstance(param, DataPointer): + param_str = f"{param.name}: DataPointer" + elif isinstance(param, EnvStream): + # env stream is not part of the FFI function signature + # continue to skip append + continue + else: + raise TypeError(f"Unsupported parameter type: {type(param)}") + + param_strs.append(param_str) + + return f"{name}({', '.join(param_strs)})" + + +def create_map_tensor_dtype_f4x2_to_f4_spec(f4_tensor_spec: Tensor) -> Tensor: + """ + Create a new Tensor spec that can be translated to f4 via f4x2->f4 conversion. + + Parameters + ---------- + f4_tensor_spec : Tensor + The original Tensor spec that declares a f4 tensor. + + Returns + ------- + Tensor + The new Tensor spec that can be translated to f4 via f4x2->f4 conversion. + """ + if f4_tensor_spec.dtype != tvm_ffi.dtype("float4_e2m1fn"): + raise ValueError("f4_tensor_spec must be a float4 tensor") + + def find_stride_one_index() -> int: + if f4_tensor_spec.strides is None: + return len(f4_tensor_spec.shape) - 1 + for i, stride in enumerate(f4_tensor_spec.strides): + if isinstance(stride, int) and stride == 1: + return i + raise ValueError("Cannot find dimension with stride=1") + stride_one_index = find_stride_one_index() + + def divisibility_divide_by_2(value: Var) -> Optional[int]: + if value.divisibility is not None: + if value.divisibility % 2 != 0: + raise ValueError(f"Dimension with stride=1 must be divisible by 2") + return value.divisibility // 2 + return None + + def map_shape(index: int, value: Union[int, Var]) -> Union[int, Var]: + if index == stride_one_index: + if isinstance(value, int): + if value % 2 != 0: + raise ValueError(f"Dimension {index} with stride=1 must be even") + return value // 2 + # create a new var with the same name and dtype + return Var(value.name, value.dtype, divisibility=divisibility_divide_by_2(value)) + return value + + def map_stride(index: int, value: Union[int, Var]) -> Union[int, Var]: + if index != stride_one_index: + if isinstance(value, int): + if value % 2 != 0: + raise ValueError(f"Dimension {index} with stride != 1 must be even") + return value // 2 + # create a new var with the same name and dtype + return Var(value.name, value.dtype, divisibility=divisibility_divide_by_2(value)) + return value + + new_shape = [map_shape(i, x) for i, x in enumerate(f4_tensor_spec.shape)] + if f4_tensor_spec.strides is not None: + new_strides = [map_stride(i, x) for i, x in enumerate(f4_tensor_spec.strides)] + else: + new_strides = None + + return Tensor( + f4_tensor_spec.name, + new_shape, + dtype=tvm_ffi.dtype("float4_e2m1fnx2"), + strides=new_strides, + map_tensor_dtype_f4x2_to_f4=True, + ) 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 new file mode 100644 index 00000000..0fe414f9 --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/tvm_ffi_builder.py @@ -0,0 +1,1781 @@ +# 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 +# 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. + +"""TVM-FFI builder for MLIR code generation.""" + +from collections.abc import Sequence +from enum import IntEnum +from typing import Callable, Optional, Union + +try: + import tvm_ffi +except ModuleNotFoundError: + pass + +from . import spec +from ..._mlir import ir +from ..._mlir.dialects import llvm +from .mlir_builder import MLIRBuilder +from dataclasses import dataclass + + +@dataclass +class CallContext: + """Call context that contains the information of the call.""" + + # the function name + fn_name: str + # the module + module: ir.Module + # the function operation + # can be used to insert llvm.alloca at beginning via + # `with ir.InsertionPoint(entry_block.operations[0]):` + entry_block: ir.Block + # the parameters of the call + params: list[spec.Param] + # the current working stream through environment synced + # with stream set in caller framework's stream context + # this is queried by the tensor device type and id + # useful for APIs where stream is not explicitly passed in + env_stream: Optional[ir.Value] + # the matched var binding + matched_var_binding: dict[spec.Var, ir.Value] + # raw arguments + raw_args: list[ir.Value] + # arg index + raw_num_args: ir.Value + # result + raw_result: ir.Value + # Keep a handle to the TVMFFIFunctionBuilder when building this context, + # so call providers can access the builder to emit globals + # (strings, helpers) in the same module as the call context. + builder: Optional["TVMFFIFunctionBuilder"] = None + + +class CallProvider: + """Call provider that implements a specific calling convention.""" + + def __call__(self, current_block: ir.Block, context: CallContext) -> ir.Block: + """Call the call provider. + + Parameters + ---------- + current_block: ir.Block + The current block to emit the call. + + context: CallContext + The call context that contains the related information of the call. + + Returns + ------- + ir.Block + The new updated current block if any. + """ + raise NotImplementedError("Call provider not implemented") + + +class TVMFFITypeIndex(IntEnum): + """TVM-FFI type index, follow the same style as C naming for now.""" + + kTVMFFINone = 0 + kTVMFFIInt = 1 + kTVMFFIBool = 2 + kTVMFFIFloat = 3 + kTVMFFIOpaquePtr = 4 + kTVMFFIDataType = 5 + kTVMFFIDevice = 6 + kTVMFFIDLTensorPtr = 7 + kTVMFFIRawStr = 8 + kTVMFFIByteArrayPtr = 9 + kTVMFFIObjectRValueRef = 10 + kTVMFFISmallStr = 11 + kTVMFFISmallBytes = 12 + kTVMFFIStaticObjectBegin = 64 + kTVMFFIObject = 64 + kTVMFFIStr = 65 + kTVMFFIBytes = 66 + kTVMFFIError = 67 + kTVMFFIFunction = 68 + kTVMFFIShape = 69 + kTVMFFITensor = 70 + kTVMFFIArray = 71 + kTVMFFIMap = 72 + kTVMFFIModule = 73 + kTVMFFIOpaquePyObject = 74 + + +class TVMFFIBuilder(MLIRBuilder): + """Base builder that provides common data structure related manipulations.""" + + def __init__(self) -> None: + super().__init__() + # this is a number we can tune to minimize the register size + # it is 6 by default to minimize the register size + self.set_raised_from_cstr_parts_max_num_parts = 6 + self.set_raised_from_cstr_parts_cache: dict[int, str] = {} + self.tvm_ffi_any_type = self.struct_type( + name="TVMFFIAny", + fields=[ + # type_index: i32 + self.i32_type, + # padding: i32 + self.i32_type, + # v_int64: i64 + self.i64_type, + ], + ) + self.tvm_ffi_object_type = self.struct_type( + name="TVMFFIObject", + fields=[ + # combined_ref_count: i64 + self.i64_type, + # type_index: i32 + self.i32_type, + # padding: i32 + self.i32_type, + # deleter: i64 (use i64 to ensure 64-bit alignment) + self.i64_type, + # cell: i64 (this is a placeholder so we can get ptr of cell) + self.i64_type, + ], + ) + self.dl_device_type = self.struct_type( + name="DLDevice", + fields=[ + # device_type: i32 (DLDeviceType enum) + self.i32_type, + # device_id: i32 + self.i32_type, + ], + ) + # DLDataType: {code: i8, bits: i8, lanes: i16} + self.dl_data_type = self.struct_type( + name="DLDataType", + fields=[ + # code: i8 (DLDataTypeCode enum) + self.i8_type, + # bits: i8 + self.i8_type, + # lanes: i16 + self.i16_type, + ], + ) + self.dltensor_type = self.struct_type( + name="DLTensor", + fields=[ + # 0 - data: void* + self.ptr_type, + # 1 - device: DLDevice + self.dl_device_type, + # 2 - ndim: i32 + self.i32_type, + # 3 - dtype: DLDataType + self.dl_data_type, + # 4 - shape: i64* + self.ptr_type, + # 5 - strides: i64* + self.ptr_type, + # 6 - byte_offset: i64 + self.i64_type, + ], + ) + self.tvm_ffi_shape_cell_type = self.struct_type( + name="TVMFFIShapeCell", + fields=[ + # shape: i64* + self.ptr_type, + # size: usize -- we use ptr type as it is same as usize + # we can load ptr then use ptrtoint to get usize + self.ptr_type, + ], + ) + self.tvm_ffi_array_cell_type = self.struct_type( + name="TVMFFIArrayCell", + fields=[ + # data: void* + self.ptr_type, + # size: i64 + self.i64_type, + ], + ) + self.tvm_ffi_func_type = self.func_type( + ret=self.i32_type, + params=[self.ptr_type, self.ptr_type, self.i32_type, self.ptr_type], + ) + + def get_object_cell_ptr(self, obj: ir.Value) -> ir.Value: + """Get the cell from the tvm_ffi_object struct. + + Parameters + ---------- + obj : ir.Value + The object ptr + + Returns + ------- + ir.Value + The pointer to the cell following the object header. + + """ + return self.getelementptr( + obj, + [0, 4], + elem_type=self.tvm_ffi_object_type, + ) + + def load_ffi_any_array_item_type_index(self, args: ir.Value, index: int) -> ir.Value: + """Get the type index from the index-th field of tvm_ffi_any_type struct. + + Semantics as follows: + + .. code-block:: c + int32_t get_type_index(void* args, const int index) { + return ((TVMFFIAny*)args)[index].type_index; + } + + Parameters + ---------- + args : ir.Value + The args pointer. + index : int + The index of the args array. + + Returns + ------- + ir.Value + The type index. + + """ + # Single GEP: args[index].type_index + # Indices: [index, 0] = access index-th element, then first field + type_index_ptr = self.getelementptr( + args, + [index, 0], + elem_type=self.tvm_ffi_any_type, + ) + return llvm.load(self.i32_type, type_index_ptr) + + def load_ffi_any_array_item_v_int64( + self, + args: ir.Value, + index: int + ) -> ir.Value: + """Get the v_int64 from the index-th field of tvm_ffi_any_type struct. + + Semantics as follows: + + .. code-block:: c + + int64_t get_v_int64(void* args, const int index) { + return ((TVMFFIAny*)args)[index].v_int64; + } + """ + v_int64_ptr = self.getelementptr( + args, + [index, 2], + elem_type=self.tvm_ffi_any_type, + ) + return llvm.load(self.i64_type, v_int64_ptr) + + def load_ffi_any_array_item_v_float64(self, args: ir.Value, index: int) -> ir.Value: + """Get the v_float64 from the index-th field of tvm_ffi_any_type struct. + + Semantics as follows: + + .. code-block:: c + + double get_v_float64(void* args, const int index) { + return ((TVMFFIAny*)args)[index].v_float64; + } + """ + v_float64_ptr = self.getelementptr( + args, + [index, 2], + elem_type=self.tvm_ffi_any_type, + ) + return llvm.load(self.f64_type, v_float64_ptr) + + def load_ffi_any_array_item_v_ptr( + self, + args: ir.Value, + index: int, + address_space: Optional[int] = None + ) -> ir.Value: + """Get the v_ptr from the index-th field of tvm_ffi_any_type struct. + + Semantics as follows: + + .. code-block:: c + + void* get_v_ptr(void* args, const int index) { + return ((TVMFFIAny*)args)[index].v_ptr; + } + """ + v_ptr_ptr = self.getelementptr( + args, + [index, 2], + elem_type=self.tvm_ffi_any_type, + ) + ptr_type = self.ptr_type_with_address_space(address_space) + return llvm.load(ptr_type, v_ptr_ptr) + + def load_shape_cell_data_ptr(self, shape_cell: ir.Value) -> ir.Value: + """Get the data pointer from the shape cell.""" + data_ptr = self.getelementptr( + shape_cell, + [0, 0], + elem_type=self.tvm_ffi_shape_cell_type, + ) + return llvm.load(self.ptr_type, data_ptr) + + def load_shape_cell_size_as_i64(self, shape_cell: ir.Value) -> ir.Value: + """Get the size from the shape cell as i64.""" + size_ptr = self.getelementptr( + shape_cell, + [0, 1], + elem_type=self.tvm_ffi_shape_cell_type, + ) + size_as_ptr_type = llvm.load(self.ptr_type, size_ptr) + return llvm.ptrtoint(self.i64_type, size_as_ptr_type) + + def load_array_cell_data_ptr(self, array_cell: ir.Value) -> ir.Value: + """Get the data pointer from the array cell.""" + data_ptr = self.getelementptr( + array_cell, + [0, 0], + elem_type=self.tvm_ffi_array_cell_type, + ) + return llvm.load(self.ptr_type, data_ptr) + + def load_array_cell_size_as_i64(self, array_cell: ir.Value) -> ir.Value: + """Get the size from the array cell as i64.""" + size_ptr = self.getelementptr( + array_cell, + [0, 1], + elem_type=self.tvm_ffi_array_cell_type, + ) + return llvm.load(self.i64_type, size_ptr) + + def load_i64_array_item(self, data: ir.Value, index: int) -> ir.Value: + """Load a shape value at the given index from the shape pointer.""" + # Get pointer to the specific shape element at index + shape_elem_ptr = self.getelementptr( + data, + [index], + elem_type=self.i64_type, + ) + # Load the actual strides value + return llvm.load(self.i64_type, shape_elem_ptr) + + def load_dltensor_data_ptr(self, dltensor: ir.Value) -> ir.Value: + """Get the data pointer from the DLTensor struct.""" + # Get pointer to the data field (first field at index 0) + data_ptr = self.getelementptr( + dltensor, + [0, 0], + elem_type=self.dltensor_type, + ) + # Load the actual data pointer + return llvm.load(self.ptr_type, data_ptr) + + def load_dltensor_device_type(self, dltensor: ir.Value) -> ir.Value: + """Get the device type from the DLTensor struct.""" + # Get pointer to the device_type field (device at index 1, then device_type at index 0) + device_type_ptr = self.getelementptr( + dltensor, + [0, 1, 0], + elem_type=self.dltensor_type, + ) + # Load the actual device type value + return llvm.load(self.i32_type, device_type_ptr) + + def load_dltensor_device_id(self, dltensor: ir.Value) -> ir.Value: + """Get the device id from the DLTensor struct.""" + # Get pointer to the device_id field (device field at index 1, then device_id at index 1) + device_id_ptr = self.getelementptr( + dltensor, + [0, 1, 1], + elem_type=self.dltensor_type, + ) + # Load the actual device id value + return llvm.load(self.i32_type, device_id_ptr) + + def load_dltensor_dtype_code(self, dltensor: ir.Value) -> ir.Value: + """Get the dtype code from the DLTensor struct.""" + # Get pointer to the dtype code field (dtype field at index 3, then code at index 0) + dtype_code_ptr = self.getelementptr( + dltensor, + [0, 3, 0], + elem_type=self.dltensor_type, + ) + # Load the actual dtype code value + return llvm.load(self.i8_type, dtype_code_ptr) + + def load_dltensor_dtype_bits(self, dltensor: ir.Value) -> ir.Value: + """Get the dtype bits from the DLTensor struct.""" + # Get pointer to the dtype bits field (dtype field at index 3, then bits at index 1) + dtype_bits_ptr = self.getelementptr( + dltensor, + [0, 3, 1], + elem_type=self.dltensor_type, + ) + # Load the actual dtype bits value + return llvm.load(self.i8_type, dtype_bits_ptr) + + def load_dltensor_dtype_lanes(self, dltensor: ir.Value) -> ir.Value: + """Get the dtype lanes from the DLTensor struct.""" + # Get pointer to the dtype lanes field (dtype field at index 3, then lanes at index 2) + dtype_lanes_ptr = self.getelementptr( + dltensor, + [0, 3, 2], + elem_type=self.dltensor_type, + ) + # Load the actual dtype lanes value + return llvm.load(self.i16_type, dtype_lanes_ptr) + + def load_dltensor_ndim(self, dltensor: ir.Value) -> ir.Value: + """Get the number of dimensions from the DLTensor struct.""" + # Get pointer to the ndim field (third field at index 2) + ndim_ptr = self.getelementptr( + dltensor, + [0, 2], + elem_type=self.dltensor_type, + ) + # Load the actual ndim value + return llvm.load(self.i32_type, ndim_ptr) + + def load_dltensor_shape(self, dltensor: ir.Value) -> ir.Value: + """Get the shape value at the given index from the DLTensor struct.""" + # Get pointer to the shape array (fifth field at index 4) + shape_ptr = self.getelementptr( + dltensor, + [0, 4], + elem_type=self.dltensor_type, + ) + return llvm.load(self.ptr_type, shape_ptr) + + def load_dltensor_strides(self, dltensor: ir.Value) -> ir.Value: + """Get the strides value at the given index from the DLTensor struct.""" + # Get pointer to the strides array (sixth field at index 5) + strides_ptr = self.getelementptr( + dltensor, + [0, 5], + elem_type=self.dltensor_type, + ) + return llvm.load(self.ptr_type, strides_ptr) + + def load_dltensor_byte_offset(self, dltensor: ir.Value) -> ir.Value: + """Get the byte offset from the DLTensor struct.""" + # Get pointer to the byte_offset field (seventh field at index 6) + byte_offset_ptr = self.getelementptr( + dltensor, + [0, 6], + elem_type=self.dltensor_type, + ) + # Load the actual byte offset value + return llvm.load(self.i64_type, byte_offset_ptr) + + def downcast_i64_to_lower_bits( + self, v_int64: ir.Value, target_dtype: "tvm_ffi.dtype" + ) -> ir.Value: + """Downcast i64 to lower bits.""" + overflow_flags = llvm.IntegerOverflowFlags.none + if target_dtype.bits == 64: + result = v_int64 + elif target_dtype.bits == 32: + result = llvm.trunc( + res=self.i32_type, arg=v_int64, overflow_flags=overflow_flags + ) + elif target_dtype.bits == 16: + result = llvm.trunc( + res=self.i16_type, arg=v_int64, overflow_flags=overflow_flags + ) + elif target_dtype.bits == 8: + result = llvm.trunc( + res=self.i8_type, arg=v_int64, overflow_flags=overflow_flags + ) + elif target_dtype.bits == 1: + # For i1 (boolean), convert i64 to boolean by checking if non-zero + result = llvm.icmp(llvm.ICmpPredicate.ne, v_int64, self.i64(0)) + else: + raise ValueError(f"Unsupported Var dtype: {target_dtype}") + return result + + def is_contiguous( + self, + expected_shape: list[Union[spec.Var, int]], + loaded_shape: list[ir.Value], + loaded_strides: list[ir.Value], + ) -> ir.Value: + """Check if the DLTensor is contiguous.""" + expected_stride: Union[int, ir.Value] = 1 + cond: ir.Value = self.i1(1) + + for index in reversed(range(len(loaded_shape))): + # strides[i] == expected_stride or shape[i] == 1 + # still constant + if isinstance(expected_stride, int): + stride_cond = self.equal( + loaded_strides[index], self.i64(expected_stride) + ) + else: + stride_cond = self.equal(loaded_strides[index], expected_stride) + + cond = self.and_( + cond, + self.or_( + stride_cond, + self.equal(loaded_shape[index], self.i64(1)), + ), + ) + if index > 0: + # try to stay in constant compute as much as possible + if isinstance(expected_shape[index], int) and isinstance( + expected_stride, int + ): + expected_stride = ( + expected_shape[index] * expected_stride # type: ignore[operator] + ) + else: + if isinstance(expected_stride, int): + expected_stride = self.i64(expected_stride) + # Handle case where expected_shape[index] might be a Var + if isinstance(expected_shape[index], int): + shape_as_i64 = self.i64(expected_shape[index]) # type: ignore[arg-type] + expected_stride = self.mul(shape_as_i64, expected_stride) + else: + # expected_shape[index] is a spec.Var, use loaded_shape[index] + expected_stride = self.mul(loaded_shape[index], expected_stride) + return cond + + def get_or_create_set_raised_from_cstr_parts(self, num_parts: int) -> str: + r"""Get or create a helper function to call TVMFFIErrorSetRaisedFromCStrParts. + + The expected generated function is as follows: + + .. code-block:: c + + void __tvm_ffi__set_error_from_parts_n( + const char* kind, + int32_t num_actual_parts, + const char* p0, + const char* p1, + ... + const char* pN-1 + ) { + const char* message_parts[n]; + message_parts[0] = p0; + message_parts[1] = p1; + ... + message_parts[n-1] = pN-1; + TVMFFIErrorSetRaisedFromCStrParts(kind, message_parts, num_actual_parts); + } + + Parameters + ---------- + num_parts : int + The number of string parts needed. + + Returns + ------- + str + The name of the helper function. + + """ + # Check cache first + if num_parts in self.set_raised_from_cstr_parts_cache: + return self.set_raised_from_cstr_parts_cache[num_parts] + + helper_name = f"__tvm_ffi__set_error_from_parts_{num_parts}" + + # Check if function already exists in the module + if self.find_func_in_module(self.module, helper_name): + self.set_raised_from_cstr_parts_cache[num_parts] = helper_name + return helper_name + + # Build the parameter list: kind, num_actual_parts, p0, p1, ..., pN-1 + param_types = [self.ptr_type] # kind + param_types.append(self.i32_type) # num_actual_parts + for _ in range(num_parts): + param_types.append(self.ptr_type) # p0, p1, ..., pN-1 + + # Create the helper function + with ir.InsertionPoint(self.module.body): # type: ignore[union-attr] + params, entry_block = self.function( + name=helper_name, + params_type=param_types, + ret_type=self.void_type, + internal=True, + ) + + kind_param = params[0] + num_actual_parts_param = params[1] + part_params = params[2:] + + with ir.InsertionPoint(entry_block): + # Allocate array of pointers to hold the message parts + message_parts_array = llvm.alloca( + res=self.ptr_type, + elem_type=self.ptr_type, + array_size=self.i32(num_parts), + alignment=8, + ) + + # Store each part in the array + for i, part_param in enumerate(part_params): + part_ptr = self.getelementptr( + message_parts_array, + [i], + elem_type=self.ptr_type, + ) + llvm.store(value=part_param, addr=part_ptr) + + # Call TVMFFIErrorSetRaisedFromCStrParts(kind, message_parts, num_actual_parts) + llvm.call( + result=None, + callee="TVMFFIErrorSetRaisedFromCStrParts", + callee_operands=[ + kind_param, + message_parts_array, + num_actual_parts_param, + ], + op_bundle_sizes=[], + op_bundle_operands=[], + ) + + # Return void + self.return_() + + self.set_raised_from_cstr_parts_cache[num_parts] = helper_name + return helper_name + + def raise_error_and_return( + self, error_kind: str, error_message_parts: list[Union[str, ir.Value]] + ) -> None: + """Raise an error and return -1. + + Instead of concatenating parts at compile time, we define each part as a global string + and call a helper function that passes them to TVMFFIErrorSetRaisedFromCStrParts. + This allows better string deduplication across the codebase. + """ + error_kind_symbol = self.define_global_string(content=error_kind) + + # Calculate actual_num_parts (max(self.set_raised_from_cstr_parts_max_num_parts, num_parts)) + call_num_parts = max( + self.set_raised_from_cstr_parts_max_num_parts, len(error_message_parts) + ) + + # Get or create the helper function for this number of parts + helper_name = self.get_or_create_set_raised_from_cstr_parts( + num_parts=call_num_parts + ) + + # Build the call operands: kind, num_actual_parts, p0, p1, ..., pN-1, (nulls...) + call_operands = [self.address_of(error_kind_symbol, self.ptr_type)] + call_operands.append(self.i32(len(error_message_parts))) # num_actual_parts + + # Add non-null part pointers + # Define global strings for each part, or forward ir.Values directly + for part in error_message_parts: + if isinstance(part, str): + part_symbol = self.define_global_string(content=part) + call_operands.append(self.address_of(part_symbol, self.ptr_type)) + else: + call_operands.append(part) + + # Pad with null pointers for unused slots + # Create null pointer constant once and reuse it + if call_num_parts > len(error_message_parts): + null_ptr = llvm.inttoptr(self.ptr_type, self.i64(0)) + for _ in range(call_num_parts - len(error_message_parts)): + call_operands.append(null_ptr) + + # Call the helper function + llvm.call( + result=None, + callee=helper_name, + callee_operands=call_operands, + op_bundle_sizes=[], + op_bundle_operands=[], + ) + + self.return_(self.i32(-1)) + + def check_condition( + self, + current_block: ir.Block, + fcond: Callable[[], ir.Value], + error_kind: str, + error_message_parts: list[str], + ) -> ir.Block: + """Check a condition and throw an error if false. + + Parameters + ---------- + current_block : ir.Block + The current block. + fcond : Callable[[], ir.Value] + Function that returns the condition to check. + error_kind : str + The kind of the error. + error_message_parts : list[str] + The message of the error. + + Returns + ------- + ir.Block + The continuation block. + + """ + error_block = current_block.create_after() + subsequent_block = error_block.create_after() + with ir.InsertionPoint(current_block): + self.cond_br( + cond=fcond(), + true_block=subsequent_block, + false_block=error_block, + # likely to be true + branch_weights=self.BRANCH_WEIGHTS_LIKELY + ) + with ir.InsertionPoint(error_block): + self.raise_error_and_return(error_kind, error_message_parts) + return subsequent_block + + +class TVMFFIFunctionBuilder(TVMFFIBuilder): + """Builder that contains specific logic for function parameters decoding.""" + + module: ir.Module + current_fn_signature: str + _fn_call_context: str + matched_var_binding: dict[spec.Var, ir.Value] + matched_var_source: dict[spec.Var, ir.Value] + + def __init__(self, module: ir.Module) -> None: + super().__init__() + self.module = module + self.current_fn_signature: str = "" + self._fn_call_context: str = "" + self.matched_var_binding = {} + self.matched_var_source = {} + + def decode_param_int( + self, current_block: ir.Block, param: spec.Var, args: ir.Value, arg_index: int + ) -> ir.Block: + """Decode the integer parameter at the given index.""" + # read the type index + with ir.InsertionPoint(current_block): + type_index: ir.Value = self.load_ffi_any_array_item_type_index(args, arg_index) + # Check if type is int or bool (both use v_int64, bool can be converted to int) + is_int = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIInt)) + is_bool = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIBool)) + is_int_or_bool = self.or_(is_int, is_bool) + + # Break error message into reusable parts for better string deduplication + current_block = self.check_condition( + current_block, + lambda: is_int_or_bool, + "TypeError", + [ + "Mismatched type on argument ", + f"#{arg_index}", + self._fn_call_context, + ", expected int", + ], + ) + with ir.InsertionPoint(current_block): + v_int64: ir.Value = self.load_ffi_any_array_item_v_int64(args, arg_index) + if param.dtype.lanes != 1: + raise ValueError(f"Unsupported Var dtype: {param.dtype}") + + return self.set_or_check_matched_var_binding( + current_block, + param, + v_int64, + [ + "value on argument ", + f"#{arg_index}", + self._fn_call_context, + ], + ) + + def decode_param_float( + self, current_block: ir.Block, param: spec.Var, args: ir.Value, arg_index: int + ) -> ir.Block: + """Decode the float parameter at the given index.""" + # read the type index + with ir.InsertionPoint(current_block): + type_index: ir.Value = self.load_ffi_any_array_item_type_index(args, arg_index) + # Check if type is float, int, or bool (int and bool can be converted to float) + is_float = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIFloat)) + is_int = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIInt)) + is_bool = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIBool)) + is_int_or_bool = self.or_(is_int, is_bool) + + if param.dtype.lanes != 1: + raise ValueError(f"Unsupported Var dtype: {param.dtype}") + + # Determine result type + if param.dtype.bits == 64: + result_type = self.f64_type + elif param.dtype.bits == 32: + result_type = self.f32_type + else: + raise ValueError(f"Unsupported Var dtype: {param.dtype}") + + # Create all blocks in the control flow + error_block = current_block.create_after() + float_block = error_block.create_after() + int_bool_check_block = float_block.create_after() + int_bool_block = int_bool_check_block.create_after() + result_block = int_bool_block.create_after() + + # Add block arguments to the result block + result_block.add_argument(result_type, ir.Location.unknown()) + + # Branch directly to float or int/bool check + with ir.InsertionPoint(current_block): + # First check if it's float + self.cond_br( + cond=is_float, + true_block=float_block, + false_block=int_bool_check_block, + ) + + # Handle float type + with ir.InsertionPoint(float_block): + v_float64: ir.Value = self.load_ffi_any_array_item_v_float64(args, arg_index) + if param.dtype.bits == 64: + float_result = v_float64 + elif param.dtype.bits == 32: + float_result = llvm.fptrunc(res=self.f32_type, arg=v_float64) + else: + raise ValueError(f"Unsupported Var dtype: {param.dtype}") + self.br(result_block, args=[float_result]) + + # In int/bool check block, verify it's actually int or bool + with ir.InsertionPoint(int_bool_check_block): + self.cond_br( + cond=is_int_or_bool, + true_block=int_bool_block, + false_block=error_block, + branch_weights=self.BRANCH_WEIGHTS_LIKELY, + ) + + # Handle int or bool type (convert to float) + with ir.InsertionPoint(int_bool_block): + v_int64: ir.Value = self.load_ffi_any_array_item_v_int64(args, arg_index) + # Convert int64 to float64 first, then to target type + v_float64_from_int = llvm.sitofp(res=self.f64_type, arg=v_int64) + if param.dtype.bits == 64: + int_bool_result = v_float64_from_int + elif param.dtype.bits == 32: + int_bool_result = llvm.fptrunc( + res=self.f32_type, arg=v_float64_from_int + ) + else: + raise ValueError(f"Unsupported Var dtype: {param.dtype}") + self.br(result_block, args=[int_bool_result]) + + # Error block + with ir.InsertionPoint(error_block): + # Break error message into reusable parts for better string deduplication + self.raise_error_and_return( + "TypeError", + [ + "Mismatched type on argument ", + f"#{arg_index}", + self._fn_call_context, + ", expected float", + ], + ) + + # Merge the results using block argument + with ir.InsertionPoint(result_block): + result = result_block.arguments[0] + self.matched_var_binding[param] = result + self.matched_var_source[param] = v_float64 + + return result_block + + def decode_param_opaque_handle( + self, + current_block: ir.Block, + param: spec.Var, + args: ir.Value, + arg_index: int, + *, + allow_int_as_ptr: bool = False, + address_space: Optional[int] = None, + ) -> ir.Block: + """Decode the opaque handle parameter at the given index.""" + # read the type index + with ir.InsertionPoint(current_block): + type_index: ir.Value = self.load_ffi_any_array_item_type_index(args, arg_index) + # Check if type is opaque pointer + is_opaque_ptr = self.equal( + type_index, self.i32(TVMFFITypeIndex.kTVMFFIOpaquePtr) + ) + # Check if type is a nullptr + is_nullptr = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFINone)) + is_opaque_ptr_or_nullptr = self.or_(is_opaque_ptr, is_nullptr) + if allow_int_as_ptr: + is_int = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIInt)) + is_opaque_ptr_or_nullptr = self.or_(is_opaque_ptr_or_nullptr, is_int) + + expect_message = ", expected handle" + if allow_int_as_ptr: + expect_message += " or int" + + # Break error message into reusable parts for better string deduplication + current_block = self.check_condition( + current_block, + lambda: is_opaque_ptr_or_nullptr, + "TypeError", + [ + "Mismatched type on argument ", + f"#{arg_index}", + self._fn_call_context, + expect_message, + ], + ) + + with ir.InsertionPoint(current_block): + # Load the opaque handle (v_ptr field contains the void*) + v_ptr: ir.Value = self.load_ffi_any_array_item_v_ptr( + args, arg_index, address_space=address_space + ) + # For opaque handles, we store the pointer directly + self.matched_var_binding[param] = v_ptr + self.matched_var_source[param] = v_ptr + + return current_block + + def check_int_value_dtype_bound( + self, + current_block: ir.Block, + value: ir.Value, + dtype: "tvm_ffi.dtype", + error_msg_context: list[str], + ) -> ir.Block: + """Check if the value is within the bounds.""" + if dtype.bits == 64: + # skip check for 64-bit integers + return current_block + + from tvm_ffi._dtype import DataTypeCode + is_uint = dtype.type_code == DataTypeCode.UINT + + # compute out the upper and lower bounds + if is_uint: + lower_bound = 0 + upper_bound = (1 << dtype.bits) - 1 + else: + lower_bound = -(1 << (dtype.bits - 1)) + upper_bound = (1 << (dtype.bits - 1)) - 1 + + error_msg = [ + "Out of bound ", + *error_msg_context, + f", expected to be in {str(dtype)} range [{lower_bound}, {upper_bound}]" + ] + + # Check bounds using appropriate comparison predicates + with ir.InsertionPoint(current_block): + lower_bound_value = self.i64(lower_bound) + upper_bound_value = self.i64(upper_bound) + if is_uint: + is_above_lower = llvm.icmp(llvm.ICmpPredicate.uge, value, lower_bound_value) + is_below_upper = llvm.icmp(llvm.ICmpPredicate.ule, value, upper_bound_value) + else: + is_above_lower = llvm.icmp(llvm.ICmpPredicate.sge, value, lower_bound_value) + is_below_upper = llvm.icmp(llvm.ICmpPredicate.sle, value, upper_bound_value) + + in_bounds = self.and_(is_above_lower, is_below_upper) + + return self.check_condition( + current_block, + lambda: in_bounds, + "ValueError", + error_msg + ) + + def check_int_value_divisibility( + self, + current_block: ir.Block, + value: ir.Value, + divisibility: int, + error_msg_context: list[str], + *, + skip_check_predicate: Optional[ir.Value] = None, + ) -> ir.Block: + """Check if the value is divisible by the specified divisibility. + + Parameters + ---------- + current_block : ir.Block + The current block to insert checks into. + value : ir.Value + The i64 value to check. + divisibility : int + The divisibility constraint. + error_msg_context : list[str] + Context for error messages. + skip_check_predicate : Optional[ir.Value], optional + The predicate to skip checking + + Returns + ------- + ir.Block + The updated block after checks. + """ + def check_divisibility() -> ir.Value: + cond = self.i64_divisible_const(value, divisibility) + if skip_check_predicate is not None: + cond = self.or_(skip_check_predicate, cond) + return cond + + error_msg = [ + "Invalid ", + *error_msg_context, + f", expected to be divisible by {divisibility}" + ] + + return self.check_condition( + current_block, + check_divisibility, + "ValueError", + error_msg + ) + + def set_or_check_matched_var_binding( + self, + current_block: ir.Block, + var: Union[spec.Var, int], + value: ir.Value, + error_msg_context: list[str], + *, + skip_check_predicate: Optional[ir.Value] = None, + ) -> ir.Block: + """Set or check the matched var binding.""" + error_kind = "ValueError" + expected_value: ir.Value + error_prefix_mismatch = "Mismatched " + + if isinstance(var, spec.Var): + # if var contains llvm_value and is not populated, populate it + if var not in self.matched_var_binding: + current_block = self.check_int_value_dtype_bound( + current_block, value, var.dtype, error_msg_context + ) + # check divisibility if specified + if var.divisibility is not None: + current_block = self.check_int_value_divisibility( + current_block, value, var.divisibility, error_msg_context, + skip_check_predicate=skip_check_predicate, + ) + # store the source value with parameter info + with ir.InsertionPoint(current_block): + self.matched_var_source[var] = value + self.matched_var_binding[var] = self.downcast_i64_to_lower_bits( + value, var.dtype + ) + return current_block + # otherwise, it appears more than once, we need to check if the value matches + expected_value = self.matched_var_source[var] + error_msg_mismatch = [ + error_prefix_mismatch, + *error_msg_context, + ", symbolic constraint violated" + ] + else: + assert isinstance(var, int) + with ir.InsertionPoint(current_block): + expected_value = self.i64(var) + error_msg_mismatch = [ + error_prefix_mismatch, + *error_msg_context, + f", expected to be {var}" + ] + return self.check_condition( + current_block, + lambda: self.equal(value, expected_value), + error_kind, + error_msg_mismatch + ) + + def set_or_check_matched_var_binding_from_shape( + self, + current_block: ir.Block, + var: Union[spec.Var, int], + value: ir.Value, + field: str, + arg_index: int, + shape_index: int, + *, + skip_check_predicate: Optional[ir.Value] = None, + ) -> ir.Block: + """Load the shape value from the argument or match the shape value from the parameter.""" + error_msg = [ + field, + f"[{shape_index}] on argument ", + f"#{arg_index}", + self._fn_call_context, + ] + return self.set_or_check_matched_var_binding( + current_block, var, value, error_msg, skip_check_predicate=skip_check_predicate + ) + + def decode_param_shape_from_ffi_array( + self, current_block: ir.Block, param: spec.Shape, arg_index: int, array_cell: ir.Value + ) -> tuple[ir.Block, list[ir.Value]]: + """Decode the shape parameter from the TVMFFIArrayCell.""" + with ir.InsertionPoint(current_block): + array_data = self.load_array_cell_data_ptr(array_cell) + array_size = self.load_array_cell_size_as_i64(array_cell) + + # Check that the array size matches the expected shape size + current_block = self.check_condition( + current_block, + lambda: self.equal(array_size, self.i64(len(param.shape))), + "ValueError", + [ + "Mismatched Shape on argument ", + f"#{arg_index}", + self._fn_call_context, + f", expected shape size={len(param.shape)}", + ], + ) + + # Load and validate each element of the array + load_shapes = [] + for i in range(len(param.shape)): + with ir.InsertionPoint(current_block): + type_index: ir.Value = self.load_ffi_any_array_item_type_index(array_data, i) + # Check if type is int or bool (both use v_int64, bool can be converted to int) + is_int = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIInt)) + + # Check that the element is an integer + current_block = self.check_condition( + current_block, + lambda: is_int, + "TypeError", + [ + f"Invalid shape element type ", + f"{param.name}[{i}]", + f" on argument ", + f"#{arg_index}", + self._fn_call_context, + f", expected int", + ], + ) + + with ir.InsertionPoint(current_block): + v_int64: ir.Value = self.load_ffi_any_array_item_v_int64(array_data, i) + load_shapes.append(v_int64) + + return (current_block, load_shapes) + + def decode_param_shape_from_ffi_shape( + self, current_block: ir.Block, param: spec.Shape, arg_index: int, shape_cell: ir.Value + ) -> tuple[ir.Block, list[ir.Value]]: + """Decode the shape parameter from the TVMFFIShapeCell.""" + with ir.InsertionPoint(current_block): + shape_data = self.load_shape_cell_data_ptr(shape_cell) + shape_size = self.load_shape_cell_size_as_i64(shape_cell) + + # Break error message into reusable parts for better string deduplication + current_block = self.check_condition( + current_block, + lambda: self.equal(shape_size, self.i64(len(param.shape))), + "ValueError", + [ + "Mismatched Shape on argument ", + f"#{arg_index}", + self._fn_call_context, + f", expected shape size={len(param.shape)}", + ], + ) + + with ir.InsertionPoint(current_block): + load_shapes = [self.load_i64_array_item(shape_data, i) for i in range(len(param.shape))] + + return (current_block, load_shapes) + + def decode_param_shape( + self, current_block: ir.Block, param: spec.Shape, args: ir.Value, arg_index: int + ) -> ir.Block: + """Decode the shape parameter at the given index.""" + with ir.InsertionPoint(current_block): + type_index: ir.Value = self.load_ffi_any_array_item_type_index(args, arg_index) + # Check if type is ffi.Shape or ffi.Array + is_ffi_shape = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIShape)) + is_ffi_array = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIArray)) + + # Create error block and subsequent blocks + error_block = current_block.create_after() + ffi_shape_block = error_block.create_after() + ffi_array_check_block = ffi_shape_block.create_after() + ffi_array_block = ffi_array_check_block.create_after() + # Create subsequent_block with i64 arguments for each shape dimension + subsequent_block = ffi_array_block.create_after(*[self.i64_type] * len(param.shape)) + + # Branch from current_block: check FFI shape first + with ir.InsertionPoint(current_block): + self.cond_br(is_ffi_shape, ffi_shape_block, ffi_array_check_block) + + # ffi shape block + with ir.InsertionPoint(ffi_shape_block): + shape_cell = self.get_object_cell_ptr( + self.load_ffi_any_array_item_v_ptr(args, arg_index) + ) + ffi_shape_block, load_shapes = self.decode_param_shape_from_ffi_shape( + ffi_shape_block, param, arg_index, shape_cell + ) + with ir.InsertionPoint(ffi_shape_block): + self.br(subsequent_block, args=load_shapes) + + # ffi array check block: verify it's actually an Array + with ir.InsertionPoint(ffi_array_check_block): + self.cond_br( + is_ffi_array, ffi_array_block, error_block, + branch_weights=self.BRANCH_WEIGHTS_LIKELY + ) + + # ffi array block + with ir.InsertionPoint(ffi_array_block): + array_cell_ptr = self.get_object_cell_ptr( + self.load_ffi_any_array_item_v_ptr(args, arg_index) + ) + ffi_array_block, load_shapes = self.decode_param_shape_from_ffi_array( + ffi_array_block, param, arg_index, array_cell_ptr + ) + with ir.InsertionPoint(ffi_array_block): + self.br(subsequent_block, args=load_shapes) + + # error block + with ir.InsertionPoint(error_block): + # Break error message into reusable parts for better string deduplication + self.raise_error_and_return( + "TypeError", + [ + "Mismatched type on argument ", + f"#{arg_index}", + self._fn_call_context, + ", expected ffi.Shape or ffi.Array", + ], + ) + + # Set or check the matched variable bindings for each dimension + with ir.InsertionPoint(subsequent_block): + shape_values = list(subsequent_block.arguments) + for i, dim in enumerate(param.shape): + subsequent_block = self.set_or_check_matched_var_binding_from_shape( + subsequent_block, dim, shape_values[i], f"{param.name}", arg_index, i + ) + + return subsequent_block + + def decode_param_tensor_dltensor_ptr( + self, + current_block: ir.Block, + param: spec.Tensor, + args: ir.Value, + arg_index: int, + ) -> tuple[ir.Block, ir.Value]: + """Decode tensor step0: check index and find out DLTensor*.""" + # read the type index + with ir.InsertionPoint(current_block): + type_index: ir.Value = self.load_ffi_any_array_item_type_index(args, arg_index) + is_ffi_tensor = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFITensor)) + is_dl_tensor = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIDLTensorPtr)) + + # Create error block and subsequent block + error_block = current_block.create_after() + ffi_tensor_block = error_block.create_after() + dl_tensor_check_block = ffi_tensor_block.create_after() + dl_tensor_block = dl_tensor_check_block.create_after() + subsequent_block = dl_tensor_block.create_after(self.ptr_type) + + # Branch from current_block: check FFI tensor first + with ir.InsertionPoint(current_block): + self.cond_br(is_ffi_tensor, ffi_tensor_block, dl_tensor_check_block) + + # ffi tensor block + with ir.InsertionPoint(ffi_tensor_block): + ffi_tensor_ptr: ir.Value = self.load_ffi_any_array_item_v_ptr(args, arg_index) + dl_tensor_ptr = self.get_object_cell_ptr(ffi_tensor_ptr) + self.br(subsequent_block, args=[dl_tensor_ptr]) + + # dltensor check block: verify it's actually a DLTensor + with ir.InsertionPoint(dl_tensor_check_block): + self.cond_br( + is_dl_tensor, dl_tensor_block, error_block, + branch_weights=self.BRANCH_WEIGHTS_LIKELY + ) + + # dltensor block + with ir.InsertionPoint(dl_tensor_block): + dltensor_ptr: ir.Value = self.load_ffi_any_array_item_v_ptr(args, arg_index) + self.br(subsequent_block, args=[dltensor_ptr]) + + # error block + with ir.InsertionPoint(error_block): + # Break error message into reusable parts for better string deduplication + self.raise_error_and_return( + "TypeError", + [ + "Mismatched type on argument ", + f"#{arg_index}", + self._fn_call_context, + ", expected Tensor", + ], + ) + + # subsequent block: receive DLTensor pointer and set it to parameter + with ir.InsertionPoint(subsequent_block): + dl_tensor_ptr = subsequent_block.arguments[0] + + return (subsequent_block, dl_tensor_ptr) + + def decode_param_tensor( + self, + current_block: ir.Block, + param: spec.Tensor, + args: ir.Value, + arg_index: int, + ) -> ir.Block: + """Decode the tensor parameter at the given index.""" + current_block, dl_tensor_ptr = self.decode_param_tensor_dltensor_ptr( + current_block, param, args, arg_index + ) + with ir.InsertionPoint(current_block): + data = self.load_dltensor_data_ptr(dl_tensor_ptr) + dtype_code = self.load_dltensor_dtype_code(dl_tensor_ptr) + dtype_bits = self.load_dltensor_dtype_bits(dl_tensor_ptr) + dtype_lanes = self.load_dltensor_dtype_lanes(dl_tensor_ptr) + device_type = self.load_dltensor_device_type(dl_tensor_ptr) + device_id = self.load_dltensor_device_id(dl_tensor_ptr) + ndim = self.load_dltensor_ndim(dl_tensor_ptr) + byte_offset = self.load_dltensor_byte_offset(dl_tensor_ptr) + + # check data alignment if specified + if param.data_alignment is not None: + def check_alignment() -> ir.Value: + # Convert pointer to integer + data_as_int = llvm.ptrtoint(self.i64_type, data) + # Check if data pointer is divisible by alignment + # (uses fast path for power-of-two alignments) + return self.i64_divisible_const(data_as_int, param.data_alignment) + + current_block = self.check_condition( + current_block, + check_alignment, + "ValueError", + [ + "Misaligned Tensor data on argument ", + f"#{arg_index}", + self._fn_call_context, + f", expected data alignment={param.data_alignment} bytes", + ], + ) + + # store the matched values, these do not need constraint checks + self.matched_var_binding[param.data] = data + self.matched_var_source[param.data] = param.data + self.matched_var_binding[param.device_id] = device_id + self.matched_var_source[param.device_id] = param.device_id + # check ndim + expected_ndim = len(param.shape) + # Break error message into reusable parts for better string deduplication + current_block = self.check_condition( + current_block, + lambda: self.equal(ndim, self.i32(expected_ndim)), + "ValueError", + [ + "Mismatched Tensor on argument ", + f"#{arg_index}", + self._fn_call_context, + f", expected ndim={expected_ndim}", + ], + ) + # check device_type + # Break error message into reusable parts for better string deduplication + current_block = self.check_condition( + current_block, + lambda: self.equal(device_type, self.i32(param.dlpack_device_type)), + "ValueError", + [ + "Mismatched Tensor on argument ", + f"#{arg_index}", + self._fn_call_context, + f", expected device_type={param.device_type_name}", + ], + ) + + # check dtype + def dtype_equal() -> ir.Value: + # check dtype (code, bits, lanes) + dtype_code_match = self.equal(dtype_code, self.i8(param.dtype.type_code)) + dtype_bits_match = self.equal(dtype_bits, self.i8(param.dtype.bits)) + dtype_lanes_match = self.equal(dtype_lanes, self.i16(param.dtype.lanes)) + return self.and_( + dtype_code_match, self.and_(dtype_bits_match, dtype_lanes_match) + ) + + # Break error message into reusable parts for better string deduplication + current_block = self.check_condition( + current_block, + dtype_equal, + "ValueError", + [ + "Mismatched Tensor on argument ", + f"#{arg_index}", + self._fn_call_context, + f", expected dtype={param.dtype}", + ], + ) + # check byte_offset + # Break error message into reusable parts for better string deduplication + current_block = self.check_condition( + current_block, + lambda: self.equal(byte_offset, self.i64(0)), + "ValueError", + [ + "Mismatched Tensor on argument ", + f"#{arg_index}", + self._fn_call_context, + ", expected byte_offset=0", + ], + ) + + with ir.InsertionPoint(current_block): + shape = self.load_dltensor_shape(dl_tensor_ptr) + load_shapes = [ + self.load_i64_array_item(shape, index) for index in range(expected_ndim) + ] + strides = self.load_dltensor_strides(dl_tensor_ptr) + load_strides = [ + self.load_i64_array_item(strides, index) + for index in range(expected_ndim) + ] + + # check the shapes + for index in range(expected_ndim): + current_block = self.set_or_check_matched_var_binding_from_shape( + current_block, + param.shape[index], + load_shapes[index], + f"{param.name}.shape", + arg_index, + index, + ) + + if param.strides is not None: + for index in range(expected_ndim): + # if shape[index] == 1 then, stride value constraint does not matter + # this is specifically to avoid some corner cases where pytorch normalizes + # stride value to 1 when shape[index] == 1 + with ir.InsertionPoint(current_block): + skip_check_predicate = self.equal(load_shapes[index], self.i64(1)) + current_block = self.set_or_check_matched_var_binding_from_shape( + current_block, + param.strides[index], + load_strides[index], + f"{param.name}.strides", + arg_index, + index, + skip_check_predicate=skip_check_predicate, + ) + else: + # Break error message into reusable parts for better string deduplication + current_block = self.check_condition( + current_block, + lambda: self.is_contiguous(param.shape, load_shapes, load_strides), + "ValueError", + [ + "Mismatched Tensor on argument ", + f"#{arg_index}", + self._fn_call_context, + ", expected contiguous", + ], + ) + return current_block + + def decode_param_stream( + self, + current_block: ir.Block, + param: spec.Stream, + args: ir.Value, + arg_index: int, + ) -> ir.Block: + """Decode the stream parameter at the given index.""" + # stream is decoded as opaque handle + return self.decode_param_opaque_handle( + current_block, param.var, args, arg_index + ) + + def decode_param_data_pointer( + self, + current_block: ir.Block, + param: spec.DataPointer, + args: ir.Value, + arg_index: int, + ) -> ir.Block: + """Decode the data pointer parameter at the given index.""" + # data pointer is decoded as opaque handle + return self.decode_param_opaque_handle( + current_block, + param.var, + args, + arg_index, + allow_int_as_ptr=True, + address_space=param.address_space, + ) + + def find_env_stream(self, params: list[spec.Param]) -> Optional[ir.Value]: + """Find the working stream from the environment Tensor that is not CPU. + + Parameters + ---------- + params : list[spec.Param] + The parameters to find the working stream from. + + Returns + ------- + Optional[ir.Value] + The working stream. + """ + for param in params: + if ( + isinstance(param, spec.Tensor) + and param.dlpack_device_type != tvm_ffi.DLDeviceType.kDLCPU + ): + device_type = self.i32(param.dlpack_device_type) + device_id = self.matched_var_binding[param.device_id] + return llvm.call( + result=self.ptr_type, + callee="TVMFFIEnvGetStream", + callee_operands=[device_type, device_id], + op_bundle_sizes=[], + op_bundle_operands=[], + ) + return None + + def get_expected_num_args(self, params: list[spec.Param]) -> int: + """Get the expected number of arguments.""" + expected_num_args = 0 + for param in params: + if not isinstance(param, spec.EnvStream): + expected_num_args += 1 + return expected_num_args + + def decode_param( # noqa: PLR0911 + self, current_block: ir.Block, param: spec.Param, args: ir.Value, arg_index: int + ) -> ir.Block: + """Decode the parameter at the given index.""" + if isinstance(param, spec.Var): + if param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.INT: + return self.decode_param_int(current_block, param, args, arg_index) + elif param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.UINT: + # UINT uses the same logic as INT since both are stored in v_int64 + return self.decode_param_int(current_block, param, args, arg_index) + elif param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.FLOAT: + return self.decode_param_float(current_block, param, args, arg_index) + elif param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.HANDLE: + return self.decode_param_opaque_handle( + current_block, param, args, arg_index + ) + else: + raise ValueError(f"Unsupported parameter type: {param.dtype.type_code}") + elif isinstance(param, spec.Shape): + return self.decode_param_shape(current_block, param, args, arg_index) + elif isinstance(param, spec.Tensor): + return self.decode_param_tensor(current_block, param, args, arg_index) + elif isinstance(param, spec.Stream): + return self.decode_param_stream(current_block, param, args, arg_index) + elif isinstance(param, spec.EnvStream): + # decode of env stream is deferred after we go through all parameters + return current_block + elif isinstance(param, spec.DataPointer): + return self.decode_param_data_pointer(current_block, param, args, arg_index) + else: + raise ValueError(f"Unsupported parameter type: {type(param)}") + + def setup_env_stream_params( + self, + current_block: ir.Block, + params: list[spec.Param], + env_stream: Optional[ir.Value], + ) -> ir.Block: + """Setup the env stream parameter.""" + for param in params: + if isinstance(param, spec.EnvStream): + if env_stream is None: + raise ValueError( + f"EnvStream cannot be detected in `{self.current_fn_signature}`" + " we need parameters to contain GPU Tensors" + ) + self.matched_var_binding[param.var] = env_stream + self.matched_var_source[param.var] = env_stream + + return current_block + + def attach_ffi_func( + self, + symbol_name: str, + params: Sequence[spec.Param], + call_provider: CallProvider, + fn_display_name: Optional[str] = None, + ) -> None: + """Add a LLVM function to the current MLIR module with the given `tvm_ffi_func_name`.""" + params_list: list[spec.Param] = list(params) + # Generate the helper function to set the error from cstr parts + self.get_or_create_set_raised_from_cstr_parts( + num_parts=self.set_raised_from_cstr_parts_max_num_parts + ) + fn_display_name = ( + fn_display_name if fn_display_name is not None else symbol_name + ) + # Generate signature for error messages + self.current_fn_signature = spec.signature(fn_display_name, params_list) + self._fn_call_context = f" when calling: `{self.current_fn_signature}`" + + with ir.InsertionPoint(self.module.body): # type: ignore[union-attr] + # void TVMFFIErrorSetRaisedFromCStr( + # const char* error_kind, const char* message); + self.declare_extern_func( + "TVMFFIErrorSetRaisedFromCStr", + [self.ptr_type, self.ptr_type], + self.void_type, + ) + # void TVMFFIErrorSetRaisedFromCStrParts( + # const char* error_kind, const char* messages, int32_t num_parts); + self.declare_extern_func( + "TVMFFIErrorSetRaisedFromCStrParts", + [self.ptr_type, self.ptr_type, self.i32_type], + self.void_type, + ) + # void* TVMFFIEnvGetStream(int32_t device_type, int32_t device_id); + self.declare_extern_func( + "TVMFFIEnvGetStream", + [self.i32_type, self.i32_type], + self.ptr_type, + ) + + (handle, args, num_args, result), entry_block = self.function( + name=f"__tvm_ffi_{symbol_name}", + params_type=[ + self.ptr_type, + self.ptr_type, + self.i32_type, + self.ptr_type, + ], + ret_type=self.i32_type, + ) + expected_num_args = self.get_expected_num_args(params_list) + + # Break error message into reusable parts for better string deduplication + current_block = entry_block + current_block = self.check_condition( + current_block, + lambda: self.equal(num_args, self.i32(expected_num_args)), + "TypeError", + [ + f"Expects {expected_num_args} parameters", + self._fn_call_context, + ], + ) + + # decode parameters to populate the matched var binding + for arg_index, param in enumerate(params_list): + current_block = self.decode_param(current_block, param, args, arg_index) + + with ir.InsertionPoint(current_block): + env_stream = self.find_env_stream(params_list) + + current_block = self.setup_env_stream_params( + current_block, params_list, env_stream + ) + + # Create call context and use call provider + context = CallContext( + fn_name=symbol_name, + module=self.module, + entry_block=entry_block, + params=params_list, + env_stream=env_stream, + matched_var_binding=self.matched_var_binding, + raw_args=args, + raw_num_args=num_args, + raw_result=result, + builder=self, + ) + + # Use the call provider to process parameters + current_block = call_provider(current_block, context) + + # Return 0 (success) + with ir.InsertionPoint(current_block): + self.return_(self.i32(0)) + + +def attach_ffi_func( + module: ir.Module, + symbol_name: str, + params: Sequence[spec.Param], + call_provider: CallProvider, + fn_display_name: Optional[str] = None, +) -> None: + """Generate a TVM-FFI function with the given symbol name and call provider. + + Parameters + ---------- + module: ir.Module + The module to attach the function to. + symbol_name: str + The name of the function to attach. + params: Sequence[spec.Param] + The parameters of the function. + fn_display_name: Optional[str] = None, + The display name of the function to attach. + + call_provider: CallProvider + The call provider that implements the calling convention. + """ + builder = TVMFFIFunctionBuilder(module) + builder.attach_ffi_func(symbol_name, params, call_provider, fn_display_name) + + +def rename_tvm_ffi_function(module: ir.Module, old_name: str, new_name: str) -> None: + """Rename the TVM FFI function in the module. + + Parameters + ---------- + module: ir.Module + The module to rename the function in. + old_name: str + The old name of the function. + new_name: str + The new name of the function. + + Raises + ------ + ValueError: If the function is not found in the module. + """ + with module.context: + builder = TVMFFIFunctionBuilder(module) + fun = builder.find_func_in_module(module, f"__tvm_ffi_{old_name}") + if fun is None: + raise ValueError( + f"Function '@{f'__tvm_ffi_{old_name}'}' not found in the module." + ) + fun.attributes["sym_name"] = ir.StringAttr.get(f"__tvm_ffi_{new_name}") diff --git a/python/CuTeDSL/cutlass/base_dsl/typing.py b/python/CuTeDSL/cutlass/base_dsl/typing.py index db641811..a9db8b29 100644 --- a/python/CuTeDSL/cutlass/base_dsl/typing.py +++ b/python/CuTeDSL/cutlass/base_dsl/typing.py @@ -239,6 +239,19 @@ def get_c_pointers(obj): return [] +def arg_compatible_with_tvm_ffi(arg): + """ + Given the `arg`, check if it is a compatible argument for TVM FFI + """ + import tvm_ffi + + return ( + hasattr(arg, "__tvm_ffi_object__") + or isinstance(arg, (int, float, bool)) + or isinstance(arg, tvm_ffi.Shape) + ) + + def get_mlir_types(obj): """ Given the `obj`, recursively go through it to extract all contained MLIR types @@ -1416,6 +1429,9 @@ class Integer(Numeric, metaclass=IntegerMeta, mlir_type=T.i32, is_abstract=True) def __rxor__(self, other, *, loc=None, ip=None): return self.__xor__(other, loc=loc, ip=ip) + def __tvm_ffi_int__(self): + return self.value + class Float(Numeric, metaclass=FloatMeta, mlir_type=T.f32, is_abstract=True): """A class representing floating-point values. @@ -1496,6 +1512,9 @@ class Float(Numeric, metaclass=FloatMeta, mlir_type=T.f32, is_abstract=True): else: raise DSLRuntimeError(f"{x} to Float conversion is not supported") + def __tvm_ffi_float__(self): + return self.value + class Boolean(Integer, metaclass=IntegerMeta, width=1, signed=True, mlir_type=T.bool): """Boolean type representation in the DSL. diff --git a/python/CuTeDSL/cutlass/cute/__init__.py b/python/CuTeDSL/cutlass/cute/__init__.py index ea216b99..db42bc7d 100644 --- a/python/CuTeDSL/cutlass/cute/__init__.py +++ b/python/CuTeDSL/cutlass/cute/__init__.py @@ -10,7 +10,7 @@ # is strictly prohibited. # Use the auto-generated enum AddressSpace -from cutlass._mlir.dialects.cute import AddressSpace +from cutlass._mlir.dialects.cute import AddressSpace, CacheEvictionPriority # Explicitly import types that might be directly used by other modules. # This is a fix for using Sphinx to generate documentation @@ -29,6 +29,7 @@ from .typing import ( ComposedLayout, Pointer, Tensor, + SymInt, ) # Import everything else @@ -114,6 +115,9 @@ from .core import ( ScaledBasis, get_divisibility, Ratio, + # FastDivmod operations + FastDivmodDivisor, + fast_divmod_create_divisor, ) from .tuple import ( @@ -175,6 +179,8 @@ from .atom import ( from .algorithm import gemm, copy, basic_copy, basic_copy_if, autovec_copy, prefetch from . import arch + +from . import export from . import nvgpu from . import testing from . import runtime @@ -200,10 +206,16 @@ KeepPTX = _dsl.KeepPTX GPUArch = _dsl.GPUArch LinkLibraries = _dsl.LinkLibraries +# attach the TVM FFI ABI interface postprocessor to the DSL +from . import _tvm_ffi_args_spec_converter + +_tvm_ffi_args_spec_converter.attach_args_spec_converter() + # Explicitly export all symbols for documentation generation __all__ = [ # Core types "AddressSpace", + "CacheEvictionPriority", "Tensor", "Layout", "ComposedLayout", @@ -219,6 +231,7 @@ __all__ = [ "ThrCopy", "TensorSSA", "ReductionOp", + "SymInt", # Basic utility functions "assume", "is_integer", @@ -344,8 +357,12 @@ __all__ = [ "repeat_like", # User defined struct "struct", + # FastDivmod operations + "FastDivmodDivisor", + "fast_divmod_create_divisor", # Modules "arch", + "export", "nvgpu", "testing", "runtime", diff --git a/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py b/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py new file mode 100644 index 00000000..642f0495 --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py @@ -0,0 +1,225 @@ +# 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 +# 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.base_dsl.tvm_ffi_builder import spec +from cutlass.base_dsl.jit_executor import ExecutionArgs +from cutlass.base_dsl.common import DSLRuntimeError +from cutlass.cutlass_dsl import is_cute_algebra_type +from cutlass._mlir.dialects import cute as _cute_ir +from .runtime import _FakeStream +from .typing import Tensor, Pointer, SymInt +from .typing import ( + Numeric, + Boolean, + Int4, + Int8, + Uint8, + Int16, + Uint16, + Int32, + Uint32, + Int64, + Uint64, + Float16, + BFloat16, + Float32, + TFloat32, + Float64, + Float6E2M3FN, + Float6E3M2FN, + Float8E5M2, + Float8E4M3FN, + Float8E8M0FNU, + Float4E2M1FN, +) +import cuda.bindings.driver as cuda + +from typing import List, Dict, Any, Optional +import inspect + +NumericToTVMFFIDtype = { + Boolean: "bool", + Int4: "int4", + Int8: "int8", + Uint8: "uint8", + Int16: "int16", + Uint16: "uint16", + Int32: "int32", + Uint32: "uint32", + Int64: "int64", + Uint64: "uint64", + Float16: "float16", + BFloat16: "bfloat16", + Float32: "float32", + TFloat32: "float32", + Float64: "float64", + Float8E5M2: "float8_e5m2", + Float8E4M3FN: "float8_e4m3fn", + Float8E8M0FNU: "float8_e8m0fnu", + Float4E2M1FN: "float4_e2m1fn", + Float6E2M3FN: "float6_e2m3fn", + Float6E3M2FN: "float6_e3m2fn", +} + +AcceptableNumericTypesForScalar = [ + Boolean, + Int8, + Uint8, + Int16, + Uint16, + Int32, + Uint32, + Int64, + Uint64, + Float32, + Float64, +] + + +def _get_llvm_address_space_from_memspace( + memspace: _cute_ir.AddressSpace, +) -> Optional[int]: + if memspace == _cute_ir.AddressSpace.gmem: + return 1 + return None + + +class SymIntId: + def __init__(self, sym_int: SymInt): + self.sym_int = sym_int + + def __hash__(self): + return id(self.sym_int) + + def __eq__(self, other) -> bool: + return self.sym_int is other.sym_int + + +def _tvm_ffi_args_spec_converter( + function_name: str, + args_spec: inspect.FullArgSpec, + dynamic_args: List[Any], + dynamic_kwargs: Dict[str, Any], +): + """Convert cute algebra args to tvm ffi spec params. + + This function converts the cute arguments specs to tvm ffi spec params. + """ + exec_args = ExecutionArgs(args_spec, function_name) + rectified_args = exec_args.get_rectified_args(dynamic_args, dynamic_kwargs) + arg_names = exec_args.args_spec.args + exec_args.args_spec.kwonlyargs + + params = [] + num_dyn_shape_vars = 0 + num_dyn_stride_vars = 0 + sym_int_id_mapping = {} + + + def alloc_shape_name(): + nonlocal num_dyn_shape_vars + name = f"n{num_dyn_shape_vars}" + num_dyn_shape_vars += 1 + return name + + def alloc_stride_name(): + nonlocal num_dyn_stride_vars + name = f"s{num_dyn_stride_vars}" + num_dyn_stride_vars += 1 + return name + + def alloc_or_reuse_symint_var(value, name_alloc_func): + nonlocal sym_int_id_mapping + sym_int_id = SymIntId(value) + if sym_int_id in sym_int_id_mapping: + return sym_int_id_mapping[sym_int_id] + name = name_alloc_func() + if value.width == 32: + dtype = NumericToTVMFFIDtype[Int32] + else: + dtype = NumericToTVMFFIDtype[Int64] + var = spec.Var(name, dtype, divisibility=value.divisibility) + sym_int_id_mapping[sym_int_id] = var + return var + + for arg, arg_name in zip(rectified_args, arg_names): + arg_type = args_spec.annotations.get(arg_name, None) + if isinstance(arg, Numeric) and arg.dtype in AcceptableNumericTypesForScalar: + params.append(spec.Var(arg_name, NumericToTVMFFIDtype[arg.dtype])) + elif is_cute_algebra_type(arg_type): + shape = [] + for i in range(len(arg)): + if isinstance(arg[i], int): + shape.append(arg[i]) + elif isinstance(arg[i], SymInt): + shape.append(alloc_or_reuse_symint_var(arg[i], alloc_shape_name)) + else: + shape.append(spec.Var(alloc_shape_name(), NumericToTVMFFIDtype[arg[i].dtype])) + params.append(spec.Shape(arg_name, shape)) + elif isinstance(arg, Tensor): + shapes = [] + for i, dyn_mask in enumerate(arg.dynamic_shapes_mask): + if not dyn_mask: + shapes.append(arg.shape[i]) + elif isinstance(arg.shape[i], SymInt): + shapes.append(alloc_or_reuse_symint_var(arg.shape[i], alloc_shape_name)) + else: + shapes.append(spec.Var(alloc_shape_name(), NumericToTVMFFIDtype[Int32])) + strides = [] + + for i, dyn_mask in enumerate(arg.dynamic_strides_mask): + if not dyn_mask: + strides.append(arg.stride[i]) + elif isinstance(arg.stride[i], SymInt): + strides.append(alloc_or_reuse_symint_var(arg.stride[i], alloc_stride_name)) + else: + if hasattr(arg, "_use_32bit_stride") and arg._use_32bit_stride: + dtype = NumericToTVMFFIDtype[Int32] + else: + dtype = NumericToTVMFFIDtype[Int64] + strides.append(spec.Var(alloc_stride_name(), dtype)) + + tvm_ffi_cute_tensor = spec.Tensor( + arg_name, + shapes, + NumericToTVMFFIDtype[arg.element_type], + strides=strides, + data_alignment=arg._assumed_align, + ) + if arg.element_type == Float4E2M1FN: + tvm_ffi_cute_tensor = spec.create_map_tensor_dtype_f4x2_to_f4_spec( + tvm_ffi_cute_tensor + ) + params.append(tvm_ffi_cute_tensor) + elif isinstance(arg, Pointer): + address_space = None + if hasattr(arg, "memspace"): + address_space = _get_llvm_address_space_from_memspace(arg.memspace) + params.append(spec.DataPointer(arg_name, address_space=address_space)) + elif isinstance(arg, _FakeStream): + if arg.use_tvm_ffi_env_stream: + params.append(spec.EnvStream(arg_name)) + else: + params.append(spec.Stream(arg_name)) + elif isinstance(arg, cuda.CUstream): + params.append(spec.Stream(arg_name)) + else: + raise DSLRuntimeError(f"Unsupported argument type: {type(arg)}") + # The following code can obtain signature of the function + # that maybe useful for future debugging and usecases. + # signature = spec.signature(function_name, params) + return params + + +def attach_args_spec_converter(): + """Attach TVM FFI ABI interface postprocessor to the DSL.""" + from .. import cutlass_dsl as _dsl + + _dsl.CuTeDSL._get_dsl()._tvm_ffi_args_spec_converter = _tvm_ffi_args_spec_converter diff --git a/python/CuTeDSL/cutlass/cute/algorithm.py b/python/CuTeDSL/cutlass/cute/algorithm.py index e93e2f74..61f2c36e 100644 --- a/python/CuTeDSL/cutlass/cute/algorithm.py +++ b/python/CuTeDSL/cutlass/cute/algorithm.py @@ -221,8 +221,8 @@ def autovec_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None: upper_bound = math.gcd(upper_bound, src.iterator.max_alignment * 8) upper_bound = math.gcd(upper_bound, dst.iterator.max_alignment * 8) - # Finally, we put a cap at 128b - num_bits_per_copy = math.gcd(upper_bound, 128) + # Finally, we put a cap at 256b + num_bits_per_copy = math.gcd(upper_bound, 256) if (num_common_elements > 1) and (num_bits_per_copy % 8 == 0): num_common_elements = num_bits_per_copy // src.element_type.width @@ -355,7 +355,7 @@ def copy( .. code-block:: python - cute.copy(tma_atom, src, dst, tma_bar_ptr=mbar_ptr, mcast_mask=mask) + cute.copy(tma_atom, src, dst, tma_bar_ptr=mbar_ptr, mcast_mask=mask, cache_policy=policy) Optional predication is supported through an additional tensor parameter. For partitioned tensors with logical profile ``((ATOM_V,ATOM_REST),REST,...)``, the predication tensor must maintain profile diff --git a/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py b/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py index 869f2c11..40f4066c 100644 --- a/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py +++ b/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py @@ -10,7 +10,9 @@ # is strictly prohibited. -from cutlass.cutlass_dsl import dsl_user_op +from cutlass.base_dsl.arch import Arch +from cutlass.base_dsl.common import DSLRuntimeError +from cutlass.cutlass_dsl import CuTeDSL, dsl_user_op from cutlass._mlir import ir from cutlass._mlir.dialects import builtin, arith, llvm, vector @@ -34,15 +36,15 @@ from ..typing import ( Int8, Int32, Float16, + Float32, BFloat16, Float32, ) - @dsl_user_op def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None): """ - Convert a vector of int8 to a vector of bfloat16. + Fast conversion from int8 to bfloat16. It converts a vector of int8 to a vector of bfloat16. :param vec_i8: The input vector of int8. :type vec_i8: 1D vector of int8 @@ -51,6 +53,9 @@ def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None): :return: The output 1D vector of bfloat16 with the same length as the input vector. :rtype: 1D vector of bfloat16 """ + arch = CuTeDSL._get_dsl().get_arch_enum() + if not arch in cvt_i8_bf16_intrinsic.supported_archs: + raise DSLRuntimeError(f"cvt_i8_bf16_intrinsic is not supported on {arch}") src_pos = 0 vec_i8x4_type = ir.VectorType.get([4], Int8.mlir_type, loc=loc) vec_i8x2_type = ir.VectorType.get([2], Int8.mlir_type, loc=loc) @@ -116,7 +121,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): """ - Convert a vector of int4 to a vector of bfloat16. + Fast conversion from int4 to bfloat16. It converts a vector of int4 to a vector of bfloat16. :param vec_i4: The input vector of int4. :type vec_i4: 1D vector of int4 @@ -125,6 +130,9 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None): :return: The output 1D vector of bfloat16 with the same length as the input vector. :rtype: 1D vector of bfloat16 """ + arch = CuTeDSL._get_dsl().get_arch_enum() + if not arch in cvt_i4_bf16_intrinsic.supported_archs: + raise DSLRuntimeError(f"cvt_i4_bf16_intrinsic is not supported on {arch}") src_pos = 0 vec_i4x8_type = ir.VectorType.get([8], Int4.mlir_type, loc=loc) vec_i4x4_type = ir.VectorType.get([4], Int4.mlir_type, loc=loc) @@ -261,3 +269,18 @@ def cvt_f4e2m1_f16_intrinsic(vec_f4e2m1, length, *, loc=None, ip=None): ip=ip, ) return vec_dst + + +# Expose supported architectures via the intrinsic symbol +cvt_i8_bf16_intrinsic.supported_archs = ( + *Arch.AmpereArchs(), + *Arch.AdaArchs(), + *Arch.HopperArchs(), + *Arch.BlackwellArchs(), +) +cvt_i4_bf16_intrinsic.supported_archs = ( + Arch.sm_100a, + Arch.sm_110a, + Arch.sm_120a, + Arch.sm_121a, +) diff --git a/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py b/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py index 3b61288e..5bbfa499 100644 --- a/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py +++ b/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py @@ -10,14 +10,19 @@ # is strictly prohibited. from functools import partial -from typing import Optional, Tuple, Union, Callable +from typing import Optional, Tuple, Union, Callable, TYPE_CHECKING from typing_extensions import deprecated -from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass.cutlass_dsl import T, dsl_user_op, cutlass_arith + +import cutlass.cutlass_dsl as cutlass_dsl from cutlass._mlir import ir from cutlass._mlir.dialects import arith, llvm, nvvm, vector +if TYPE_CHECKING: + from cutlass.tensor import TensorSSA + # Forward nvvm enums from cutlass._mlir.dialects.nvvm import ( ProxyKind, @@ -27,6 +32,8 @@ from cutlass._mlir.dialects.nvvm import ( RoundingModeKind, ) +from ..core import size + from ..typing import ( Int, Boolean, @@ -184,7 +191,7 @@ def block_idx_in_cluster(*, loc=None, ip=None) -> Int32: @dsl_user_op def shuffle_sync_op( - value: Numeric, + value: Union[Numeric, "TensorSSA"], offset: Int, mask: Int = FULL_MASK, mask_and_clamp: Int = WARP_SIZE - 1, @@ -192,12 +199,12 @@ def shuffle_sync_op( *, loc=None, ip=None, -) -> Numeric: +) -> Union[Numeric, "TensorSSA"]: """ Shuffles a value within the threads of a warp. :param value: The value to shuffle - :type value: Numeric + :type value: Numeric or TensorSSA :param mask: A mask describing the threads participating in this operation :type mask: Int :param offset: A source lane or a source lane offset depending on kind @@ -211,12 +218,37 @@ def shuffle_sync_op( :return: The shuffled value :rtype: Numeric """ + from ..tensor import TensorSSA + + if isinstance(value, TensorSSA): + bit_width = value.dtype.width * size(value.shape) + if bit_width == 32: + i32_val = llvm.bitcast( + T.i32(), value.ir_value(loc=loc, ip=ip), loc=loc, ip=ip + ) + i32_res = nvvm.shfl_sync( + T.i32(), + Int32(mask).ir_value(loc=loc, ip=ip), + i32_val, + Int32(offset).ir_value(loc=loc, ip=ip), + Int32(mask_and_clamp).ir_value(loc=loc, ip=ip), + kind, + loc=loc, + ip=ip, + ) + result_vec = llvm.bitcast(value.type, i32_res, loc=loc, ip=ip) + return TensorSSA(result_vec, value.shape, value.dtype) + else: + raise ValueError(f"shuffle_sync only supports 32 bit, but got {value.type}") + if not isinstance(value, Numeric): value = as_numeric(value) + if value.width > 64: raise ValueError("shuffle_sync only supports values up to 64 bits") orig_type = type(value) + if value.width < 32: if value.dtype.is_float: value = value.to(Float32) @@ -331,6 +363,7 @@ def warp_reduction( :rtype: cutlass.Numeric """ offset = threads_in_group // 2 + while offset > 0: val = op( val, @@ -343,7 +376,8 @@ def warp_reduction( warp_reduction_max = partial( - warp_reduction, op=lambda x, y: fmax(x, y) if isinstance(x, Float32) else max(x, y) + warp_reduction, + op=lambda x, y: fmax(x, y) if isinstance(x, Float32) else cutlass_dsl.max(x, y), ) warp_reduction_sum = partial(warp_reduction, op=lambda x, y: x + y) @@ -796,11 +830,11 @@ add_packed_f32x2 = partial( calc_packed_f32x2_op, src_c=None, calc_func=nvvm.add_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(), @@ -811,7 +845,6 @@ def fmax( ) ) - @dsl_user_op def rcp_approx(a: Union[float, Float32], *, loc=None, ip=None): return Float32( diff --git a/python/CuTeDSL/cutlass/cute/atom.py b/python/CuTeDSL/cutlass/cute/atom.py index a8d05a94..3e001538 100644 --- a/python/CuTeDSL/cutlass/cute/atom.py +++ b/python/CuTeDSL/cutlass/cute/atom.py @@ -99,6 +99,9 @@ class Trait(ABC): def unpack(self, *, loc=None, ip=None, **kwargs) -> ir.Value: return self.value + def with_(self, *, loc=None, ip=None, **kwargs) -> "Trait": + return self.__class__(self.unpack(loc=loc, ip=ip, **kwargs)) + def make_atom(ty, values=None, *, loc=None, ip=None): """ @@ -185,6 +188,21 @@ class Atom(ABC): """ return self._trait.get(field, loc=loc, ip=ip) + def with_(self, *, loc=None, ip=None, **kwargs) -> "Atom": + """ + Returns a new Atom with the new Operation and Trait with the given runtime state. The runtime state + is provided as keyword arguments and it is Atom-specific. + + .. code-block:: python + + tiled_copy = cute.make_tiled_copy(tma_copy_op) + new_tiled_copy = tiled_copy.with_(tma_bar_ptr=tma_bar_ptr, cache_policy=cute.CacheEvictionPriority.EVICT_LAST) + + The ``with_`` method provides a way to the user to modify such runtime state or create an executable Atom + (e.g. an Executable TMA Load Atom). + """ + return self.__class__(self.op, self._trait.with_(loc=loc, ip=ip, **kwargs)) + def _unpack(self, *, loc=None, ip=None, **kwargs) -> ir.Value: return self._trait.unpack(loc=loc, ip=ip, **kwargs) @@ -215,23 +233,28 @@ class MmaAtom(Atom): # @property - def thr_id(self) -> Layout: - return static(self._trait.value.type.thr_id) + @dsl_user_op + def thr_id(self, *, loc=None, ip=None) -> Layout: + return static(self._trait.value.type.thr_id, loc=loc, ip=ip) @property - def shape_mnk(self) -> Shape: - return _unpack_x_tuple(self._trait.value.type.shape_mnk) + @dsl_user_op + def shape_mnk(self, *, loc=None, ip=None) -> Shape: + return _unpack_x_tuple(self._trait.value.type.shape_mnk, loc=loc, ip=ip) @property - def tv_layout_A(self) -> Layout: - return static(self._trait.value.type.layout_a_tv) + @dsl_user_op + def tv_layout_A(self, *, loc=None, ip=None) -> Layout: + return static(self._trait.value.type.layout_a_tv, loc=loc, ip=ip) @property - def tv_layout_B(self) -> Layout: - return static(self._trait.value.type.layout_b_tv) + @dsl_user_op + def tv_layout_B(self, *, loc=None, ip=None) -> Layout: + return static(self._trait.value.type.layout_b_tv, loc=loc, ip=ip) @property - def tv_layout_C(self) -> Layout: + @dsl_user_op + def tv_layout_C(self, *, loc=None, ip=None) -> Layout: return static(self._trait.value.type.layout_c_tv) # diff --git a/python/CuTeDSL/cutlass/cute/core.py b/python/CuTeDSL/cutlass/cute/core.py index 204fda58..c678f569 100644 --- a/python/CuTeDSL/cutlass/cute/core.py +++ b/python/CuTeDSL/cutlass/cute/core.py @@ -10,6 +10,7 @@ # 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 @@ -304,9 +305,10 @@ class IntValue(cutlass_arith.ArithValue): def __str__(self): if self.divisibility == 1: return "?" - else: + elif self.type.width == 32: return f"?{{div={self.divisibility}}}" - + else: + return f"?{{i{self.type.width} div={self.divisibility}}}" def __repr__(self): parent_name = cutlass_arith.ArithValue.__name__ return super().__str__().replace(parent_name, IntValue.__name__) @@ -913,8 +915,8 @@ class _Layout(Layout): layout = make_layout((4, 8), stride=(8, 1)) - # map linear index back to coordinate: 5 -> (1, 1) - coord = get_hier_coord(5, layout) + # map linear index back to coordinate: 9 -> (1, 1) + coord = layout.get_hier_coord(9) """ idx_val = Int32(idx).ir_value(loc=loc, ip=ip) crd = _cute_ir.get_hier_coord(idx_val, self, loc=loc, ip=ip) @@ -3116,12 +3118,6 @@ def make_ptr( if isinstance(value, ir.Value) and llvm.PointerType.isinstance(value.type): value = llvm.ptrtoint(T.i64(), value) - 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) @@ -3611,9 +3607,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 - - if not is_static(sliced_lay): - raise ValueError("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) @@ -4125,3 +4121,166 @@ class ThrMma(_ThrMma): @deprecated("cute.core.ThrCopy is deprecated, use cute.ThrCopy instead") class ThrCopy(_ThrCopy): pass + + +# +# FastDivmod operations for optimized division and modulus +# +class FastDivmodDivisor: + """ + First-class FastDivmod divisor with operator overloading support. + + This class wraps a FastDivmod divisor and enables natural Python operator syntax: + quotient, remainder = divmod(dividend, divisor) + quotient = dividend // divisor + remainder = dividend % divisor + + :ivar _divisor: The FastDivmod divisor MLIR value + """ + + @dsl_user_op + def __init__( + self, + divisor: Integer, + is_power_of_2: bool = None, + *, + loc=None, + ip=None, + ): + """ + Create a FastDivmod divisor for optimized division operations. + + :param divisor: The divisor value (should be runtime-dynamic value) + :param is_power_of_2: Whether divisor is known to be a power of 2. + Defaults to False. + """ + # Convert divisor to ir.Value for MLIR operation + if isinstance(divisor, ir.Value): + divisor_val = divisor + else: + divisor_val = Int32(divisor).ir_value() + + # Use user-provided flag or default to False + # Power-of-2 optimization should be handled by compiler passes at IR level + if is_power_of_2 is None: + is_power_of_2 = False + + # Create FastDivmod divisor + fast_divmod_divisor_type = _cute_ir.FastDivmodDivisorType.get(32, is_power_of_2) + self._divisor = _cute_ir.fast_divmod_create_divisor( + fast_divmod_divisor_type, divisor_val, loc=loc, ip=ip + ) + + @dsl_user_op + def __rdivmod__( + self, dividend: Integer, *, loc=None, ip=None + ) -> Tuple[Integer, Integer]: + """ + Overload for: divmod(dividend, self) + Returns (quotient, remainder). + + :param dividend: The dividend value + :param loc: Source location for MLIR + :param ip: Insertion point for MLIR + :return: Tuple of (quotient, remainder) + """ + # Convert dividend to ir.Value for MLIR operation + if isinstance(dividend, ir.Value): + dividend_val = dividend + else: + dividend_val = Int32(dividend).ir_value() + + quotient_type = dividend_val.type + remainder_type = dividend_val.type + + results = _cute_ir.fast_divmod_compute( + quotient_type, + remainder_type, + dividend_val, + self._divisor, + loc=loc, + ip=ip, + ) + return (IntValue(results[0]), IntValue(results[1])) + + @dsl_user_op + def __rfloordiv__(self, dividend: Integer, *, loc=None, ip=None) -> Integer: + """ + Overload for: dividend // self + Returns quotient only. + + :param dividend: The dividend value + :param loc: Source location for MLIR + :param ip: Insertion point for MLIR + :return: The quotient + """ + quotient, _ = self.__rdivmod__(dividend, loc=loc, ip=ip) + return quotient + + @dsl_user_op + def __rmod__(self, dividend: Integer, *, loc=None, ip=None) -> Integer: + """ + Overload for: dividend % self + Returns remainder only. + + :param dividend: The dividend value + :param loc: Source location for MLIR + :param ip: Insertion point for MLIR + :return: The remainder + """ + _, remainder = self.__rdivmod__(dividend, loc=loc, ip=ip) + return remainder + + def __extract_mlir_values__(self): + """Extract MLIR values for Host->Device transfer.""" + return [self._divisor] + + def __new_from_mlir_values__(self, values): + """Reconstruct FastDivmodDivisor from MLIR values.""" + 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 +) -> FastDivmodDivisor: + """Create a FastDivmod divisor for optimized division operations. + + This function creates a FastDivmod divisor that precomputes auxiliary values + to enable fast division and modulus operations without using division instructions. + + The returned FastDivmodDivisor object supports natural Python operator syntax: + divisor = fast_divmod_create_divisor(batch_size) + quotient, remainder = divmod(linear_idx, divisor) + quotient = linear_idx // divisor + remainder = linear_idx % divisor + + :param divisor: The divisor value (should be runtime-dynamic value) + :type divisor: Integer + :return: FastDivmodDivisor object with operator overloading support + :rtype: FastDivmodDivisor + """ + return FastDivmodDivisor(divisor, loc=loc, ip=ip) diff --git a/python/CuTeDSL/cutlass/cute/export/__init__.py b/python/CuTeDSL/cutlass/cute/export/__init__.py new file mode 100644 index 00000000..2368d211 --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/export/__init__.py @@ -0,0 +1,38 @@ +# 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 +# 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 .c_header_generator import CuteCHeaderGenerator + +from ...base_dsl.export import ( + get_export_module, + dump_to_object as _dump_to_object, + export_to_c as _export_to_c, +) +from ...cutlass_dsl import CuTeDSL +from functools import partial as _partial +from ...cutlass_dsl.cuda_jit_executor import CudaDialectJitCompiledFunction + +dump_to_object = _partial( + _dump_to_object, + dsl=CuTeDSL._get_dsl(), +) +export_to_c = _partial( + _export_to_c, + dsl=CuTeDSL._get_dsl(), + c_header_generator=CuteCHeaderGenerator(), + use_gpu_dialect=False, +) +__all__ = [ + "CuteCHeaderGenerator", + "get_export_module", + "dump_to_object", + "export_to_c", +] diff --git a/python/CuTeDSL/cutlass/cute/export/c_header_generator.py b/python/CuTeDSL/cutlass/cute/export/c_header_generator.py new file mode 100644 index 00000000..92293acf --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/export/c_header_generator.py @@ -0,0 +1,241 @@ +# 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 +# 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.cute.typing import NumericMeta, Integer +from cutlass.base_dsl.export import CHeaderGenerator +from cutlass.base_dsl.dsl import is_dynamic_expression +from cutlass.base_dsl.common import DSLRuntimeError +from cutlass.base_dsl.jit_executor import ExecutionArgs +from cutlass.cutlass_dsl.cutlass import is_cute_algebra_type +from ..runtime import Tensor, Pointer + +from typing import List, Any, Dict +from inspect import isclass +import cuda.bindings.driver as cuda + +# ============================================================================= +# Cute DSL C Header Generator for c/cpp AOT support +# ============================================================================= + + +class CuteCHeaderGenerator(CHeaderGenerator): + """This class provides a Export C Header Generator for cute c/cpp AOT support.""" + + def _get_cute_algebra_type(self, arg_type: Any, arg: Any) -> str: + """Judge if the dynamic elements of the cute algebra type are same(Int32 or Int64). + If so, generate the corresponding C type. Otherwise, refuse to generate the argument + in C header and raise an error. + """ + if not isinstance(arg, (list, tuple, Integer)): + raise DSLRuntimeError( + f"Unsupported argument for c function argument generation: {arg} with type {arg_type}" + ) + if isinstance(arg, Integer): + return self.numeric_to_c_type[arg.dtype] + dyn_type = None + for elem in arg: + if not isinstance(elem, (Integer, int)): + raise DSLRuntimeError( + f"Unsupported argument for c function argument generation: {arg} with type {arg_type}" + ) + if isinstance(elem, int): + continue + if dyn_type is not None: + if elem.dtype != dyn_type: + raise DSLRuntimeError( + f"Expects all dynamic elements of the cute algebra type to be of the same type, but got {elem.dtype} and {dyn_type}" + ) + else: + dyn_type = elem.dtype + if dyn_type is None: + return "int32_t " + return self.numeric_to_c_type[dyn_type] + + def _generate_binary_declaration(self, symbol_prefix: str): + """ + Generate the binary of the compiled function. + """ + return "" + + def _generate_kernel_metadata( + self, symbol_prefix: str, kernel_info: Dict[str, List], dsl_name: str + ): + """ + Generate the kernel metadata for the compiled function. + """ + kernel_metadata_struct = f""" +typedef struct {{ + CUlibrary module; +}} {symbol_prefix}_Kernel_Metadata_t; +""" + kernel_metadata_load = f""" +#ifdef __cplusplus +extern "C" {{ +#endif +void _mlir_{symbol_prefix}_cuda_init(void **); +void _mlir_{symbol_prefix}_cuda_load(void **); +static inline void {symbol_prefix}_Kernel_Metadata_Load({symbol_prefix}_Kernel_Metadata_t *metadata) {{ + CUlibrary *libraryPtr = &(metadata->module); + int32_t ret; + struct {{ + CUlibrary **libraryPtr; + int32_t *ret; + }} initArgs = {{&libraryPtr, &ret}}; + _mlir_{symbol_prefix}_cuda_init((void **)(&initArgs)); + {dsl_name}_CUDA_ERROR_CHECK((CUresult)(ret)); + struct {{ + CUlibrary *library; + int32_t *ret; + }} loadArgs = {{libraryPtr, &ret}}; + _mlir_{symbol_prefix}_cuda_load((void **)(&loadArgs)); + {dsl_name}_CUDA_ERROR_CHECK((CUresult)(ret)); + CUdevice device; + {dsl_name}_CUDA_ERROR_CHECK(cuCtxGetDevice(&device)); + int max_shared_memory_per_block_optin; + {dsl_name}_CUDA_ERROR_CHECK(cuDeviceGetAttribute(&max_shared_memory_per_block_optin, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, device)); + unsigned int num_kernels; + {dsl_name}_CUDA_ERROR_CHECK(cuLibraryGetKernelCount(&num_kernels, metadata->module)); + CUkernel *kernels = (CUkernel *)malloc(num_kernels * sizeof(CUkernel)); + {dsl_name}_CUDA_ERROR_CHECK(cuLibraryEnumerateKernels(kernels, num_kernels, metadata->module)); + for (unsigned int i = 0; i < num_kernels; i++) {{ + {dsl_name}_CUDA_ERROR_CHECK(cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, max_shared_memory_per_block_optin, kernels[i], device)); + }} + free(kernels); +}} +""" + kernel_metadata_unload = f""" +static inline void {symbol_prefix}_Kernel_Metadata_Unload({symbol_prefix}_Kernel_Metadata_t *metadata) {{ + {dsl_name}_CUDA_ERROR_CHECK(cuLibraryUnload(metadata->module)); +}} + +#ifdef __cplusplus +}} +#endif +""" + + return kernel_metadata_struct + kernel_metadata_load + kernel_metadata_unload + + def _generate_arguments( + self, + symbol_prefix: str, + args_spec: ExecutionArgs, + args: List[Any], + kwargs: Dict[str, Any], + ): + """ + Generate the arguments of the wrapper function. + """ + arguments = [] + packed_args = [] + declarations = [] + # traverse the runtime args_spec and generate the arguments + rectified_args = args_spec.get_rectified_args(args, kwargs) + input_arg_names = args_spec.args_spec.args + args_spec.args_spec.kwonlyargs + for arg_name, arg in zip(input_arg_names, rectified_args): + arg_type = args_spec.args_spec.annotations.get(arg_name, None) + + # process optional argument + if arg is None: + continue + + if isinstance(arg, Pointer): + arguments.append(f"void *{arg_name}") + packed_args.append("&" + arg_name) + elif isinstance(arg, Tensor): + dynamic_shapes = ( + f"\n int32_t dynamic_shapes[{sum(arg.dynamic_shapes_mask)}];" + if sum(arg.dynamic_shapes_mask) > 0 + else "" + ) + stride_type = "int32_t" if arg._use_32bit_stride else "int64_t" + dynamic_strides = ( + f"\n {stride_type} dynamic_strides[{sum(arg.dynamic_strides_mask)}];" + if sum(arg.dynamic_strides_mask) > 0 + else "" + ) + declarations.append( + f""" +typedef struct {{ + void *data;{dynamic_shapes}{dynamic_strides} +}} {symbol_prefix}_Tensor_{arg_name}_t; +""" + ) + arguments.append(f"{symbol_prefix}_Tensor_{arg_name}_t *{arg_name}") + packed_args.append(arg_name) + # Generate basic numeric types + elif isinstance(arg_type, NumericMeta): + arguments.append(self._generate_numeric_argument(arg_name, arg_type)) + packed_args.append("&" + arg_name) + elif is_cute_algebra_type(arg_type): + c_type = self._get_cute_algebra_type(arg_type, arg) + arguments.append(f"{c_type}*{arg_name}") + for i in range(self._count_dynamic_expression(arg)): + packed_args.append("&" + arg_name + "[" + str(i) + "]") + elif isclass(arg_type) and issubclass(arg_type, cuda.CUstream): + arguments.append("CUstream " + arg_name) + packed_args.append("&" + arg_name) + else: + raise DSLRuntimeError( + f"Unsupported argument for c function argument generation: {arg} with type {arg_type}" + ) + + return arguments, packed_args, declarations + + def _generate_wrapper_function( + self, + symbol_prefix: str, + args_spec: ExecutionArgs, + function_name: str, + kernel_info: Dict[str, List], + dynamic_args: list, + dynamic_kwargs: dict, + ): + """ + Generate the wrapper function for the compiled function which is provided to users as the entry point. + It uses the `symbol_prefix` as the function name for identification. The host/device symbols are hidden under the bytecode. + """ + # 1. Get the name of the function wrapper + wrapper_function_name = f"{symbol_prefix}_wrapper" + capi_function_name = f"_mlir_{symbol_prefix}__mlir_ciface_{function_name}" + # 2. Generate the signature of the wrapper function + arguments, packed_args, declarations = self._generate_arguments( + symbol_prefix, args_spec, dynamic_args, dynamic_kwargs + ) + # 3. Get the return type of the wrapper function. + # Note that this requires the return type to be properly annotated in python. + return_type = args_spec.args_spec.annotations.get("return", None) + if return_type is None: + return_type = "void" + else: + return_type = self.numeric_to_c_type[return_type][:-1] + declarations = "\n".join(declarations) + + # 4. Generate the wrapper function + function = ( + declarations + + f""" +#ifdef __cplusplus +extern "C" +#endif +void {capi_function_name}(void **args, int32_t num_args); + +static inline {return_type} {wrapper_function_name}({symbol_prefix}_Kernel_Metadata_t *metadata, {", ".join(arguments)}) {{ + {return_type} ret; + void *args[{len(packed_args) + 1}] = {{ + {", ".join(packed_args)}, + &ret + }}; + {capi_function_name}(args, {len(packed_args) + 1}); + return ret; +}} +""" + ) + return function diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py b/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py index 55f8bb86..8bd46c24 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py @@ -22,7 +22,7 @@ from cutlass._mlir import ir from ...atom import CopyOp, Trait from ...tensor import ReductionOp -from ...typing import Int16, Pointer, Integer, Numeric +from ...typing import Int16, Int64, Pointer, Integer, Numeric from ..common import OpError from ..tcgen05.mma import CtaGroup @@ -113,6 +113,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_CACHE_POLICY_FIELD_NAME = "cache_policy" class TmaCopyOp(CopyOp): @@ -184,6 +185,10 @@ class CopyBulkTensorTileG2SOp(TmaCopyOp): class CopyBulkTensorTileG2SNonExecTrait(Trait): # We allow kw args to be dropped so that the user can write common code for non-multicast # and multicast loads. + + def with_(self, *, loc=None, ip=None, **kwargs) -> "CopyBulkTensorTileG2STrait": + return CopyBulkTensorTileG2STrait(self.unpack(loc=loc, ip=ip, **kwargs)) + def unpack( self, *, @@ -191,13 +196,15 @@ class CopyBulkTensorTileG2SNonExecTrait(Trait): ip=None, tma_bar_ptr: Optional[Pointer] = None, tma_desc_ptr: Optional[Pointer] = None, + cache_policy: Optional[Int64] = None, **kwargs, ): """ Custom implementation of unpack for non-executable TMAs. The non-multicast TMA load requires a `tma_bar_ptr` keyword argument to be provided when - using `cute.copy`. Any other kw arguments will be ignored instead of triggering an error. + using `cute.copy`. `cache_policy` keyword argument to be provided to set the l2 cache eviction priority. + Any other kw arguments will be ignored instead of triggering an error. """ if not isinstance(tma_bar_ptr, Pointer): raise ValueError( @@ -215,9 +222,25 @@ class CopyBulkTensorTileG2SNonExecTrait(Trait): exec_value = _cute_nvgpu_ir.atom_set_value( exec_value, attr, tma_desc_ptr.value, loc=loc, ip=ip ) + if cache_policy is not None: + if not isinstance(cache_policy, Int64): + raise ValueError( + "expects `Int64` value to be provided via the cache_policy kw argument" + ) + + attr_str = ( + f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_CACHE_POLICY_FIELD_NAME}>" + ) + attr = ir.Attribute.parse(attr_str) + exec_value = _cute_nvgpu_ir.atom_set_value( + exec_value, attr, cache_policy.value, loc=loc, ip=ip + ) return exec_value +class CopyBulkTensorTileG2STrait(Trait): + pass + # # TMA GMEM -> SMEM multicast copies # @@ -277,6 +300,13 @@ class CopyBulkTensorTileG2SMulticastOp(TmaCopyOp): class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait): + def with_( + self, *, loc=None, ip=None, **kwargs + ) -> "CopyBulkTensorTileG2SMulticastTrait": + return CopyBulkTensorTileG2SMulticastTrait( + self.unpack(loc=loc, ip=ip, **kwargs) + ) + def unpack( self, *, @@ -285,12 +315,14 @@ class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait): tma_bar_ptr: Optional[Pointer] = None, mcast_mask=None, tma_desc_ptr=None, + cache_policy: Optional[Int64] = None, ): """ Custom implementation of unpack for non-executable TMAs. The multicast TMA load requires a `tma_bar_ptr` and a `mcast_mask` keyword arguments to be - provided when using `cute.copy`. + provided when using `cute.copy`. `cache_policy` keyword argument to be provided to set the + l2 cache eviction priority. """ if not isinstance(tma_bar_ptr, Pointer): raise ValueError( @@ -312,6 +344,19 @@ class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait): exec_value = _cute_nvgpu_ir.atom_set_value( exec_value, attr, tma_desc_ptr.value, loc=loc, ip=ip ) + if cache_policy is not None: + if not isinstance(cache_policy, Int64): + raise ValueError( + "expects `Int64` value to be provided via the cache_policy kw argument" + ) + + attr_str = ( + f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_CACHE_POLICY_FIELD_NAME}>" + ) + attr = ir.Attribute.parse(attr_str) + exec_value = _cute_nvgpu_ir.atom_set_value( + exec_value, attr, cache_policy.value, loc=loc, ip=ip + ) # Set the tma_bar_ptr at last to ensure that the atom creation and setting # operations above can be moved outside the loop attr_str = "#cute_nvgpu.atom_copy_field_tmaload" @@ -321,6 +366,9 @@ class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait): ) return exec_value +class CopyBulkTensorTileG2SMulticastTrait(Trait): + pass + # # TMA SMEM -> GMEM copies @@ -351,14 +399,24 @@ class CopyBulkTensorTileS2GOp(TmaCopyOp): def _make_trait( self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs - ) -> "CopyBulkTensorTileS2GTrait": + ) -> "CopyBulkTensorTileS2GNonExecTrait": raise NotImplementedError( "Use cpasync.make_tiled_tma_atom to obtain a copy Atom for TMA" ) -class CopyBulkTensorTileS2GTrait(Trait): - def unpack(self, *, loc=None, ip=None, tma_desc_ptr: Optional[Pointer] = None): +class CopyBulkTensorTileS2GNonExecTrait(Trait): + def with_(self, *, loc=None, ip=None, **kwargs) -> "CopyBulkTensorTileS2GTrait": + return CopyBulkTensorTileS2GTrait(self.unpack(loc=loc, ip=ip, **kwargs)) + + def unpack( + self, + *, + loc=None, + ip=None, + tma_desc_ptr: Optional[Pointer] = None, + cache_policy: Optional[Int64] = None, + ): """ Custom implementation of unpack for non-executable TMAs. """ @@ -371,9 +429,26 @@ class CopyBulkTensorTileS2GTrait(Trait): exec_value = _cute_nvgpu_ir.atom_set_value( exec_value, attr, tma_desc_ptr.value, loc=loc, ip=ip ) + if cache_policy is not None: + if not isinstance(cache_policy, Int64): + raise ValueError( + "expects `Int64` value to be provided via the cache_policy kw argument" + ) + + attr_str = ( + f"#cute_nvgpu.atom_copy_field_tmastore<{TMA_CACHE_POLICY_FIELD_NAME}>" + ) + attr = ir.Attribute.parse(attr_str) + exec_value = _cute_nvgpu_ir.atom_set_value( + exec_value, attr, cache_policy.value, loc=loc, ip=ip + ) return exec_value +class CopyBulkTensorTileS2GTrait(Trait): + pass + + @dataclass(frozen=True) class CopyReduceBulkTensorTileS2GOp(TmaCopyOp): """ @@ -400,7 +475,7 @@ class CopyReduceBulkTensorTileS2GOp(TmaCopyOp): def _make_trait( self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs - ) -> "CopyReduceBulkTensorTileS2GTrait": + ) -> "CopyReduceBulkTensorTileS2GNonExecTrait": raise NotImplementedError( "Use cpasync.make_tiled_tma_atom to obtain a copy Atom for TMA" ) @@ -426,8 +501,20 @@ class CopyReduceBulkTensorTileS2GOp(TmaCopyOp): assert False, "unrecognized self.reduction_kind" -class CopyReduceBulkTensorTileS2GTrait(Trait): - def unpack(self, *, loc=None, ip=None, tma_desc_ptr: Optional[Pointer] = None): +class CopyReduceBulkTensorTileS2GNonExecTrait(Trait): + def with_( + self, *, loc=None, ip=None, **kwargs + ) -> "CopyReduceBulkTensorTileS2GTrait": + return CopyReduceBulkTensorTileS2GTrait(self.unpack(loc=loc, ip=ip, **kwargs)) + + def unpack( + self, + *, + loc=None, + ip=None, + tma_desc_ptr: Optional[Pointer] = None, + cache_policy: Optional[Int64] = None, + ): """ Custom implementation of unpack for non-executable TMAs. """ @@ -440,9 +527,25 @@ class CopyReduceBulkTensorTileS2GTrait(Trait): exec_value = _cute_nvgpu_ir.atom_set_value( exec_value, attr, tma_desc_ptr.value, loc=loc, ip=ip ) + if cache_policy is not None: + if not isinstance(cache_policy, Int64): + raise ValueError( + "expects `Int64` value to be provided via the cache_policy kw argument" + ) + attr_str = ( + f"#cute_nvgpu.atom_copy_field_tmareduce<{TMA_CACHE_POLICY_FIELD_NAME}>" + ) + attr = ir.Attribute.parse(attr_str) + exec_value = _cute_nvgpu_ir.atom_set_value( + exec_value, attr, cache_policy.value, loc=loc, ip=ip + ) return exec_value +class CopyReduceBulkTensorTileS2GTrait(Trait): + pass + + # # Bulk GMEM -> SMEM copies # @@ -494,13 +597,15 @@ class CopyBulkG2STrait(Trait): loc=None, ip=None, mbar_ptr: Optional[Pointer] = None, + cache_policy: Optional[Int64] = None, **kwargs, ): """ Custom implementation of unpack for bulk copy load. The non-multicast bulk load requires a `mbar_ptr` keyword argument to be provided when - using `cute.copy`. Any other kw arguments will be ignored instead of triggering an error. + using `cute.copy`. `cache_policy` keyword argument to be provided to set the l2 cache eviction priority. + Any other kw arguments will be ignored instead of triggering an error. """ if not isinstance(mbar_ptr, Pointer): raise ValueError( @@ -511,6 +616,18 @@ class CopyBulkG2STrait(Trait): val = _cute_nvgpu_ir.atom_set_value( self.value, attr, mbar_ptr.value, loc=loc, ip=ip ) + if cache_policy is not None: + if not isinstance(cache_policy, Int64): + raise ValueError( + "expects `Int64` value to be provided via the cache_policy kw argument" + ) + attr_str = ( + f"#cute_nvgpu.atom_copy_field_bulkg2s<{TMA_CACHE_POLICY_FIELD_NAME}>" + ) + attr = ir.Attribute.parse(attr_str) + val = _cute_nvgpu_ir.atom_set_value( + val, attr, cache_policy.value, loc=loc, ip=ip + ) return val @@ -566,6 +683,7 @@ class CopyBulkG2SMulticastTrait(Trait): ip=None, mbar_ptr: Optional[Pointer] = None, mcast_mask: Optional[Integer] = None, + cache_policy: Optional[Int64] = None, **kwargs, ): """ @@ -592,6 +710,18 @@ class CopyBulkG2SMulticastTrait(Trait): val = _cute_nvgpu_ir.atom_set_value( val, attr, Int16(mcast_mask).ir_value(loc=loc, ip=ip), loc=loc, ip=ip ) + if cache_policy is not None: + if not isinstance(cache_policy, Int64): + raise ValueError( + "expects `Int64` value to be provided via the cache_policy kw argument" + ) + attr_str = ( + f"#cute_nvgpu.atom_copy_field_bulkg2s<{TMA_CACHE_POLICY_FIELD_NAME}>" + ) + attr = ir.Attribute.parse(attr_str) + val = _cute_nvgpu_ir.atom_set_value( + val, attr, cache_policy.value, loc=loc, ip=ip + ) return val diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py b/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py index 02fe57f8..cbbbf53b 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py @@ -35,8 +35,8 @@ from .copy import ( CopyReduceBulkTensorTileS2GOp, CopyBulkTensorTileG2SNonExecTrait, CopyBulkTensorTileG2SMulticastNonExecTrait, - CopyBulkTensorTileS2GTrait, - CopyReduceBulkTensorTileS2GTrait, + CopyBulkTensorTileS2GNonExecTrait, + CopyReduceBulkTensorTileS2GNonExecTrait, ) @@ -156,7 +156,7 @@ def make_tiled_tma_atom( loc=loc, ip=ip, ) - return atom.CopyAtom(op, CopyBulkTensorTileS2GTrait(res[0])), res[1] + return atom.CopyAtom(op, CopyBulkTensorTileS2GNonExecTrait(res[0])), res[1] elif isinstance(op, CopyReduceBulkTensorTileS2GOp): res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_reduce( gmem_tensor.value, @@ -167,7 +167,10 @@ def make_tiled_tma_atom( loc=loc, ip=ip, ) - return atom.CopyAtom(op, CopyReduceBulkTensorTileS2GTrait(res[0])), res[1] + return ( + atom.CopyAtom(op, CopyReduceBulkTensorTileS2GNonExecTrait(res[0])), + res[1], + ) else: raise ValueError(f"expects a bulk tensor (TMA) Copy Op, but got {op}") diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py index f5c2cef8..a539c5b7 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py @@ -285,7 +285,9 @@ def find_tmem_tensor_col_offset(tmem_tensor: Tensor, *, loc=None, ip=None) -> In """ tmem_col_mask = 0x0000FFFF offset = ( - core.cosize(recast_tensor(tmem_tensor, Int32).layout, loc=loc, ip=ip) + core.cosize( + recast_tensor(tmem_tensor, Int32, loc=loc, ip=ip).layout, loc=loc, ip=ip + ) & tmem_col_mask ) if isinstance(offset, int): diff --git a/python/CuTeDSL/cutlass/cute/runtime.py b/python/CuTeDSL/cutlass/cute/runtime.py index 4395d7a1..ab04f27e 100644 --- a/python/CuTeDSL/cutlass/cute/runtime.py +++ b/python/CuTeDSL/cutlass/cute/runtime.py @@ -10,19 +10,24 @@ # is strictly prohibited. import ctypes +import os +import sys +from pathlib import Path from functools import lru_cache import itertools import operator -from typing import Union, Optional +from typing import Union, Optional, Type, List # MLIR modules imports from cutlass._mlir import ir import cutlass._mlir.dialects.cute as _cute_ir +import cutlass._mlir.dialects.cuda as _cuda_dialect -from cutlass.cutlass_dsl import JitArgAdapterRegistry, DSLRuntimeError +from cutlass.cutlass_dsl import JitArgAdapterRegistry, CuTeDSL as _CuTeDSL +from cutlass.base_dsl.common import DSLRuntimeError # Local modules imports -from .typing import AddressSpace, Tensor, Type, Pointer, Numeric +from .typing import AddressSpace, Layout, Tensor, Pointer, Numeric, SymInt from . import core from .tensor import _Tensor as CoreTensor @@ -76,6 +81,9 @@ class _Pointer(Pointer): def __get_mlir_types__(self): return [self.mlir_type] + def __tvm_ffi_opaque_ptr__(self): + return self._pointer + def __c_pointers__(self): if self._c_pointer is None: self._desc = ctypes.c_void_p(int(self._pointer)) @@ -104,14 +112,6 @@ class _Pointer(Pointer): def align(self, min_align: int, *, loc=None, ip=None) -> Pointer: raise NotImplementedError("align is not supported in runtime") - def verify(self, expected_py_type): - if expected_py_type is Pointer: - return True - elif isinstance(expected_py_type, ir.Value) and expected_py_type.ty is Pointer: - return True - - return False - def __str__(self) -> str: return f"Ptr<0x{int(self._pointer):016x}@{self._addr_space}>" @@ -120,12 +120,30 @@ class _Pointer(Pointer): class _Tensor(Tensor): - def __init__(self, tensor, assumed_align=None, use_32bit_stride=False): + def __init__( + self, + tensor, + assumed_align=None, + use_32bit_stride=False, + *, + enable_tvm_ffi=False, + ): # If tensor is already a DLPack object, use it directly if hasattr(tensor, "__dlpack_device__") and not hasattr(tensor, "__dlpack__"): - self._dlpack_data = tensor + self._dlpack_data = tensor.__dlpack_device__() else: - self._dlpack_data = tensor.__dlpack__() + try: + # we expect no stream sync. Because torch has different default behavior + # for stream parameter on different version. + # we need to explicitly pass -1 to achieve no sync effects. + self._dlpack_data = tensor.__dlpack__(stream=-1) + except Exception: + self._dlpack_data = tensor.__dlpack__() + if enable_tvm_ffi: + import tvm_ffi + + self._tvm_ffi_tensor = tvm_ffi.from_dlpack(tensor) + self._dlpack_data = self._tvm_ffi_tensor.__dlpack__() self._dltensor_wrapper = None self._assumed_align = assumed_align self._is_dynamic = False @@ -217,7 +235,6 @@ class _Tensor(Tensor): mode, stride_order, divisibility ) return self - @property @lazily_load_dltensor def element_type(self) -> Type[Numeric]: @@ -342,6 +359,18 @@ class _Tensor(Tensor): def data_ptr(self): return self._dltensor_wrapper.data_ptr + @property + @lazily_load_dltensor + def dynamic_shapes_mask(self): + """Get the mask of dynamic shapes in the tensor.""" + return self._dltensor_wrapper.get_dynamic_shapes_mask() + + @property + @lazily_load_dltensor + def dynamic_strides_mask(self): + """Get the mask of dynamic strides in the tensor.""" + return self._dltensor_wrapper.get_dynamic_strides_mask() + @lazily_load_dltensor def __c_pointers__(self): self._memref_desc = self._dltensor_wrapper.build_memref_desc( @@ -357,11 +386,336 @@ class _Tensor(Tensor): assert isinstance(values[0], CoreTensor) return CoreTensor(values[0].value, self._dtype) + def __tvm_ffi_object__(self): + return self._tvm_ffi_tensor + + +def _get_cute_type_str(inp): + def _convert_dyn_elem(e): + return f"?{{i{e.width} div={e.divisibility}}}" + + elems = [_convert_dyn_elem(e) if isinstance(e, SymInt) else str(e) for e in inp] + return "(" + ",".join(elems) + ")" + + +class _FakeCompactTensor(Tensor): + def __init__( + self, + dtype, + shape, + stride_order, + memspace=None, + assumed_align=None, + use_32bit_stride=False, + ): + self._dtype = dtype + self._shape = shape + self._stride_order = stride_order or tuple(range(len(shape))) + self._memspace = memspace or AddressSpace.gmem + self._assumed_align = assumed_align or -(-dtype.width // 8) + self._use_32bit_stride = use_32bit_stride + + def __str__(self) -> str: + return f"FakeTensorOrdered<{self._dtype}, {self._shape}, {self._stride_order}>" + + def __repr__(self): + return self.__str__() + + @property + def mlir_type(self) -> ir.Type: + shape_ty = ir.Type.parse( + '!cute.shape<"' + _get_cute_type_str(self._shape) + '">' + ) + layout_ty = _cute_ir.LayoutType.get_ordered( + shape_ty, self._stride_order, self._use_32bit_stride + ) + self._stride = layout_ty.stride + ptr_ty = _cute_ir.PtrType.get( + self._dtype.mlir_type, self._memspace, self._assumed_align + ) + return _cute_ir.MemRefType.get(ptr_ty, layout_ty) + + def __get_mlir_types__(self): + return [self.mlir_type] + + def __new_from_mlir_values__(self, values): + assert len(values) == 1 + assert isinstance(values[0], CoreTensor) + return CoreTensor(values[0].value, self._dtype) + + def __setitem__(self, crd, value): + raise DSLRuntimeError("runtime._FakeCompactTensor is not indexable") + + def __getitem__(self, crd): + raise DSLRuntimeError("runtime._FakeCompactTensor is not indexable") + + @property + def element_type(self) -> Type[Numeric]: + return self._dtype + + @property + def memspace(self): + return self._memspace + + @property + def iterator(self): + raise DSLRuntimeError("runtime._FakeTensor has dummy iterator") + + @property + def shape(self): + return self._shape + + @property + def stride(self): + return self._stride + + @property + def dynamic_shapes_mask(self): + return tuple(1 if isinstance(e, SymInt) else 0 for e in self._shape) + + @property + def dynamic_strides_mask(self): + return tuple(1 if isinstance(e, SymInt) else 0 for e in self._stride) + + def fill(self, value: Numeric): + raise DSLRuntimeError("runtime._FakeCompactTensor is not writable") + + +class _FakeTensor(Tensor): + """Fake Tensor implementation as a placeholder. + It mimics the interface of Tensor, but does not hold real data or allow indexing. + Used for compilation or testing situations where only shape/type/layout information is needed. + All attempts to access or mutate data will raise errors. + """ + + """ + Create a fake tensor with the given shape, type, and layout. + + :param dtype: Data type of the tensor elements + :type dtype: Type[Numeric] + :param shape: Shape of the tensor, consists of int (static) or SymInt (dynamic) + :type shape: tuple[int, ...] + :param stride: Stride of the tensor, defaults to None, consists of int (static) or SymInt (dynamic) + :type stride: tuple[int, ...], optional + :param assumed_align: Assumed alignment of the tensor (bytes), defaults to None. If None, uses the element size bytes as the assumed alignment. + :type assumed_align: int, optional + :param use_32bit_stride: Whether to use 32-bit stride. Defaults to False. When True, the dynamic stride bitwidth + will be set to 32 for small problem sizes (cosize(layout) <= Int32_max) for better performance. This is only applied + when the dimension is dynamic. + :type use_32bit_stride: bool, optional + + """ + + def __init__(self, dtype, shape, *, stride, memspace=None, assumed_align=None): + self._dtype = dtype + self._shape = shape + self._stride = stride + self._memspace = memspace or AddressSpace.generic + self._assumed_align = assumed_align + if assumed_align is None: + # use the bytes width of the element dtype. The alignment is at least one byte align. + self._assumed_align = (self._dtype.width + 7) // 8 + + if not isinstance(shape, (tuple, list)): + raise ValueError(f"Expected tuple or list but got {type(shape)}") + + if not all(isinstance(s, (int, SymInt)) for s in self._shape): + raise ValueError("All shape elements must be int or SymInt") + + if stride is not None and not all( + isinstance(s, (int, SymInt)) for s in self._stride + ): + raise ValueError("All stride elements must be int or SymInt") + @property + def mlir_type(self) -> ir.Type: + shape_str = _get_cute_type_str(self._shape) + stride_str = _get_cute_type_str(self._stride) + layout_ty = ir.Type.parse(f'!cute.layout<"{shape_str}:{stride_str}">') + ptr_ty = _cute_ir.PtrType.get( + self._dtype.mlir_type, self._memspace, self._assumed_align + ) + return _cute_ir.MemRefType.get(ptr_ty, layout_ty) + + def __get_mlir_types__(self): + return [self.mlir_type] + + def __new_from_mlir_values__(self, values): + assert len(values) == 1 + assert isinstance(values[0], CoreTensor) + return CoreTensor(values[0].value, self._dtype) + + def __str__(self) -> str: + return f"FakeTensor<{self._dtype}, {self._shape}, {self._stride}>" + + def __repr__(self): + return self.__str__() + + def __setitem__(self, crd, value): + raise DSLRuntimeError("runtime._FakeTensor is not indexable") + + def __getitem__(self, crd): + raise DSLRuntimeError("runtime._FakeTensor is not indexable") + + @property + def element_type(self) -> Type[Numeric]: + return self._dtype + + @property + def memspace(self): + return self._memspace + + @property + def iterator(self): + raise DSLRuntimeError("runtime._FakeTensor has dummy iterator") + + @property + def shape(self): + return self._shape + + @property + def stride(self): + return self._stride + + @property + def dynamic_shapes_mask(self): + return tuple(1 if isinstance(e, SymInt) else 0 for e in self._shape) + + @property + def dynamic_strides_mask(self): + return tuple(1 if isinstance(e, SymInt) else 0 for e in self._stride) + + def fill(self, value: Numeric): + raise DSLRuntimeError("runtime._FakeTensor is not writable") + + +def make_fake_compact_tensor( + dtype, + shape, + *, + stride_order=None, + memspace=None, + assumed_align=None, + use_32bit_stride=False, +): + """ + Create a fake tensor with the specified shape, element type, and a compact memory layout. + + :param dtype: Data type of the tensor elements. + :type dtype: Type[Numeric] + :param shape: Shape of the tensor. + :type shape: tuple[int, ...] + :param stride_order: Order in which strides (memory layout) are assigned to the tensor dimensions. + If None, the default layout is row-major. Otherwise, it should be a permutation of the dimension indices. + :type stride_order: tuple[int, ...], optional + :param memspace: Memory space where the fake tensor resides. Optional. + :type memspace: str, optional + :param assumed_align: Assumed byte alignment for the tensor data. If None, the default alignment is used. + :type assumed_align: int, optional + :param use_32bit_stride: Whether to use 32-bit stride for dynamic dimensions. If True and the total size of the + layout (cosize(layout)) fits within int32, then dynamic strides will use 32-bit integers for improved performance. + Only applies when dimensions are dynamic. Defaults to False. + :type use_32bit_stride: bool, optional + :return: An instance of a fake tensor with the given properties and compact layout. + :rtype: _FakeCompactTensor + + **Examples:** + + .. code-block:: python + + @cute.jit + def foo(x: cute.Tensor): + ... + + x = make_fake_compact_tensor( + cutlass.Float32, (100, cute.sym_int32(divisibility=8)), stride_order=(1, 0) + ) + + # Compiled function will take a tensor with the type: + # tensor o (100,?{div=8}):(?{i32 div=8},1)> + compiled_foo = cute.compile(foo, x) + """ + + return _FakeCompactTensor( + dtype, + shape, + stride_order=stride_order, + memspace=memspace, + assumed_align=assumed_align, + use_32bit_stride=use_32bit_stride, + ) + +def make_fake_tensor(dtype, shape, stride, *, memspace=None, assumed_align=None): + """ + Create a fake tensor with the specified element type, shape, and stride. + + :param dtype: Data type of the tensor elements. + :type dtype: Type[Numeric] + :param shape: Shape of the tensor. + :type shape: tuple[int, ...] + :param stride: Stride of the tensor. + :type stride: tuple[int, ...] + :param assumed_align: Assumed byte alignment for the tensor data. If None, the default alignment is used. Defaults to None. + :type assumed_align: int, optional + :return: An instance of a fake tensor with the given properties. + :rtype: _FakeTensor + """ + return _FakeTensor( + dtype, shape, stride=stride, memspace=memspace, assumed_align=assumed_align + ) + + +class _FakeStream: + """A fake stream that can be used as a placeholder for a stream in compilation. + + When use_tvm_ffi_env_stream is True and the function is compiled with TVM-FFI, + the argument will be skipped from the function signature and we pass in + this value through the environment stream obtained from caller context + (e.g. torch.cuda.current_stream()). + """ + + use_tvm_ffi_env_stream: bool + + def __init__(self, *, use_tvm_ffi_env_stream: bool = False): + self.use_tvm_ffi_env_stream = use_tvm_ffi_env_stream + + def __str__(self) -> str: + return f"FakeStream" + + def __repr__(self): + return self.__str__() + + def __new_from_mlir_values__(self, values): + assert len(values) == 1 + return values[0] + + def __c_pointers__(self): + return [0] + + def __get_mlir_types__(self): + return [_cuda_dialect.StreamType.get()] + + +def make_fake_stream(*, use_tvm_ffi_env_stream: bool = False): + """Create a fake stream that can be used as a placeholder for a stream in compilation. + + When use_tvm_ffi_env_stream is True and the function is compiled with TVM-FFI, + the argument will be skipped from the function signature and we pass in + this value through the environment stream obtained from caller context + (e.g. torch.cuda.current_stream()). This can speedup the calling process + since we no longer need to do stream query in python. + + :param use_tvm_ffi_env_stream: Whether to skip this parameter use environment stream instead. + :type use_tvm_ffi_env_stream: bool + """ + return _FakeStream(use_tvm_ffi_env_stream=use_tvm_ffi_env_stream) + def from_dlpack( tensor_dlpack, assumed_align=None, use_32bit_stride=False, + *, + enable_tvm_ffi=False, ) -> Tensor: """Convert from tensor object supporting __dlpack__() to a CuTe Tensor. @@ -374,25 +728,33 @@ def from_dlpack( stride bitwidth will be set to 32 for small problem size (cosize(layout) <= Int32_max) for better performance. This is only applied when the dimension is dynamic. :type use_32bit_stride: bool, optional + :param enable_tvm_ffi: Whether to enable TVM-FFI, defaults to False. When True, the tensor will be converted to + a TVM-FFI function compatible tensor. + :type enable_tvm_ffi: bool, optional :return: A CuTe Tensor object :rtype: Tensor - Examples: - .. code-block:: python + **Examples:** - import torch - from cutlass.cute.runtime import from_dlpack - x = torch.randn(100, 100) - y = from_dlpack(x) - y.shape - # (100, 100) - type(y) - # + .. code-block:: python + + import torch + from cutlass.cute.runtime import from_dlpack + x = torch.randn(100, 100) + y = from_dlpack(x) + y.shape + # (100, 100) + type(y) + # """ + # If the environment variable `CUTE_DSL_ENABLE_TVM_FFI` is set to True, the tensor will be converted to + # a TVM-FFI function compatible tensor. + enable_tvm_ffi = enable_tvm_ffi or _CuTeDSL._get_dsl().envar.enable_tvm_ffi return _Tensor( tensor_dlpack, assumed_align=assumed_align, use_32bit_stride=use_32bit_stride, + enable_tvm_ffi=enable_tvm_ffi, ) @@ -451,6 +813,23 @@ def make_ptr( return _Pointer(address_value, dtype, mem_space, assumed_align=assumed_align) +def nullptr( + dtype: Type[Numeric], + mem_space: AddressSpace = AddressSpace.generic, + assumed_align=None, +) -> Pointer: + """Create a null pointer which is useful for compilation + + :param dtype: Data type of the pointer elements + :type dtype: Type[Numeric] + :param mem_space: Memory address space, defaults to AddressSpace.generic + :type mem_space: AddressSpace, optional + :return: A null pointer object + :rtype: Pointer + """ + return make_ptr(dtype, 0, mem_space, assumed_align=assumed_align) + + class TensorAdapter: """ Convert a DLPack protocol supported tensor/array to a cute tensor. @@ -469,6 +848,74 @@ class TensorAdapter: return self._arg.__get_mlir_types__() +def find_runtime_libraries(*, enable_tvm_ffi: bool = True) -> List[str]: + """ + Find the runtime libraries that needs to be available for loading modules. + + :param enable_tvm_ffi: Whether to enable TVM-FFI. + :type enable_tvm_ffi: bool, optional + :return: A list of runtime libraries that needs to be available for loading modules. + :rtype: list + """ + + def _get_cuda_dialect_runtime_path(): + libs = os.environ.get("CUTE_DSL_LIBS") + if libs: + sep = ";" if sys.platform.startswith("win32") else ":" + for path in libs.split(sep): + if path.endswith("libcuda_dialect_runtime.so"): + return path + try: + # find package library from wheel package + pkg_base = Path(__file__).resolve().parent.parent + lib_path = pkg_base / "lib" / "libcuda_dialect_runtime.so" + if lib_path.is_file(): + return str(lib_path) + except OSError: + return None + + libs = [] + cuda_dialect_runtime_path = _get_cuda_dialect_runtime_path() + if cuda_dialect_runtime_path: + libs.append(cuda_dialect_runtime_path) + + if enable_tvm_ffi: + import tvm_ffi + + libs.append(tvm_ffi.libinfo.find_libtvm_ffi()) + + return libs + +# cache to load runtime libraries so they can be found by the DSO loader +_LOAD_MODULE_LIBS_CACHE = [] + + +def load_module(file_path: str, *, enable_tvm_ffi: bool = True): + """Load a module from a file path. Today only support TVM-FFI module. + + :param file_path: The path to the module file + :type file_path: str + :param enable_tvm_ffi: Whether to enable TVM-FFI, defaults to True. When True, the module will be loaded as a TVM-FFI module. + :type enable_tvm_ffi: bool, optional + :return: A module object + :rtype: module + """ + if len(_LOAD_MODULE_LIBS_CACHE) == 0: + # ensure the runtime libraries are loaded so they can be found by the DSO loader + # no need to load tvm_ffi library here since it will be loaded by tvm_ffi package. + for path in find_runtime_libraries(enable_tvm_ffi=False): + if Path(path).exists(): + _LOAD_MODULE_LIBS_CACHE.append(ctypes.CDLL(path)) + + if enable_tvm_ffi: + import tvm_ffi + + return tvm_ffi.load_module(file_path) + else: + raise DSLRuntimeError( + "Unimplemented, please load the module with enable_tvm_ffi=True." + ) + # ------------------------------------------------------------------------- # Try to register_jit_arg_adapter for TensorAdapter # ------------------------------------------------------------------------- @@ -479,7 +926,6 @@ try: # Register for numpy.ndarray JitArgAdapterRegistry.register_jit_arg_adapter(numpy.ndarray)(TensorAdapter) except ImportError: pass # silent attempt, suppress error - try: # Register for torch.Tensor import torch diff --git a/python/CuTeDSL/cutlass/cute/tensor.py b/python/CuTeDSL/cutlass/cute/tensor.py index 96834f2b..d28c5713 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, + CuTeDSL, ) from cutlass._mlir import ir import cutlass._mlir.dialects.cute as _cute_ir @@ -33,6 +34,7 @@ from .core import ( _pack_coord, _pack_shape, _ComposedLayout, + _ComposedLayoutWithInnerFunc, is_static, is_weakly_congruent, rank, @@ -250,7 +252,7 @@ class _Tensor(Tensor): val = tensor[0] # Error: sub-byte scalar dereference not supported """ if has_underscore(crd): - return slice_(self.value, crd, loc=loc, ip=ip) + return slice_(self, crd, loc=loc, ip=ip) elif isinstance(self.type, _cute_ir.CoordTensorType): res = _cute_ir.get_iter( slice_(self, crd, loc=loc, ip=ip).value, loc=loc, ip=ip @@ -448,7 +450,7 @@ class _Tensor(Tensor): self._check_can_load_store() - res_vect = _cute_ir.memref_load_vec(self.value, row_major=True, loc=loc, ip=ip) + res_vect = _cute_ir.memref_load_vec(self.value, 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}" @@ -501,9 +503,7 @@ 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, row_major=True, loc=loc, ip=ip - ) + return _cute_ir.memref_store_vec(new_data, self.value, loc=loc, ip=ip) @dsl_user_op def fill(self, value: Numeric, *, loc=None, ip=None) -> None: @@ -555,6 +555,11 @@ class _Tensor(Tensor): ): raise ValueError(f"{self} doesn't support load and store") + if self.type.is_swizzled: + raise NotImplementedError( + f"load & store swizzled memory is not supported yet: {self}" + ) + def _check_can_dereference(self): # Check for sub-byte types and raise error if needed if self.element_type.width % 8 != 0 and self.element_type is not Boolean: @@ -579,11 +584,16 @@ def make_tensor( evaluates coordinates by applying the layout mapping and dereferencing the engine at the resulting offset. - :param iterator: Engine component (pointer, iterator, or counting iterator) that provides - data access capabilities - :type iterator: Union[Pointer, IntTuple] + :param iterator: Engine component that provides data access capabilities. Can be: + - A pointer (Pointer type) + - An integer or integer tuple for coordinate tensors + - A shared memory descriptor (SmemDescType) + :type iterator: Union[Pointer, IntTuple, ir.Value] :param layout: Layout component that defines the mapping from logical coordinates to - physical offsets + physical offsets. Can be: + - A shape tuple that will be converted to a layout + - A Layout object + - A ComposedLayout object (must be a normal layout) :type layout: Union[Shape, Layout, ComposedLayout] :param loc: Source location for MLIR operation tracking, defaults to None :type loc: Optional[Location] @@ -592,30 +602,40 @@ def make_tensor( :return: A tensor object representing the composition E ∘ L :rtype: Tensor - :raises ValueError: If iterator type is not supported + :raises TypeError: If iterator type is not a supported type + :raises ValueError: If layout is a composed layout with customized inner functions **Examples:** .. code-block:: python - # Create a tensor with row-major layout + # Create a tensor with row-major layout from a pointer + ptr = make_ptr(Float32, base_ptr, AddressSpace.gmem) layout = make_layout((64, 128), stride=(128, 1)) tensor = make_tensor(ptr, layout) - # Create a tensor with hierarchical layout + # Create a tensor with hierarchical layout in shared memory + smem_ptr = make_ptr(Float16, base_ptr, AddressSpace.smem) layout = make_layout(((128, 8), (1, 4, 1)), stride=((32, 1), (0, 8, 4096))) tensor = make_tensor(smem_ptr, layout) - # Create a coord tensor + # Create a coordinate tensor layout = make_layout(2, stride=16 * E(0)) - tensor = make_tensor(5, layout) + tensor = make_tensor(5, layout) # coordinate tensor with iterator starting at 5 Notes: - The engine (iterator) must support random access operations - Common engine types include raw pointers, arrays, and random-access iterators - The layout defines both the shape (logical dimensions) and stride (physical mapping) - Supports both direct coordinate evaluation T(c) and partial evaluation (slicing) + - ComposedLayouts must be "normal" layouts (no inner functions) + - For coordinate tensors, the iterator is converted to a counting sequence """ + if isinstance(layout, _ComposedLayoutWithInnerFunc): + raise ValueError( + "CuTe DSL tensor does not support composed layouts with inner functions: {layout}" + ) + if not isinstance(layout, (Layout, ComposedLayout)): layout = make_layout(layout, loc=loc, ip=ip) # Automatic decay to normal layout @@ -964,6 +984,65 @@ def print_tensor( _cute_ir.print_view(tensor.value, verbose=verbose, is_signed=signed, loc=loc, ip=ip) +def _get_row_and_col_map(col_maj_shape_1d: tuple, is_row_to_col: bool): + """ + Create an index mapping mask for converting between row-major and column-major vector ordering. + """ + + # create row-major layout with compact row-major stride + if col_maj_shape_1d is None: + raise ValueError("vector shape for row/col map cannot be None") + if isinstance(col_maj_shape_1d, tuple): + row_maj_shape_1d = tuple(reversed(flatten(col_maj_shape_1d))) + else: + # Single dimension + row_maj_shape_1d = col_maj_shape_1d + + flat_shape = flatten(row_maj_shape_1d) + if isinstance(flat_shape, tuple): + strides = [] + current_stride = 1 + flat_shape = tuple(reversed(flat_shape)) + for dim in flat_shape: + strides.append(current_stride) + current_stride *= dim + row_maj_stride = tuple(reversed(strides)) + else: + # Single dimension + row_maj_stride = 1 + + row_maj_lay_1d = make_layout(row_maj_shape_1d, stride=row_maj_stride) + + # get idx map + shape_size = size(row_maj_shape_1d) + mask = [0] * shape_size + + for col_index in range(shape_size): + row_index = crd2idx(col_index, row_maj_lay_1d) + if is_row_to_col: + mask[row_index] = col_index + else: + mask[col_index] = row_index + + return mask + + +def _row2col(vec: ir.Value, *, shape, loc=None, ip=None) -> ir.Value: + """ + Convert a vector or tensor from row-major order to column-major order. + """ + row_and_col_map = _get_row_and_col_map(shape, is_row_to_col=True) + return vector.shuffle(vec, vec, row_and_col_map, loc=loc, ip=ip) + + +def _col2row(vec: ir.Value, *, shape, loc=None, ip=None) -> ir.Value: + """ + Convert a vector or tensor from column-major order to row-major order. + """ + row_and_col_map = _get_row_and_col_map(shape, is_row_to_col=False) + return vector.shuffle(vec, vec, row_and_col_map, loc=loc, ip=ip) + + def _infer_broadcast_shape(*shapes: Shape) -> Shape: """ Infer the broadcasted shape from multiple input shapes according to broadcasting rules. @@ -1129,6 +1208,32 @@ class TensorSSA(cutlass_arith.ArithValue): return res + @dsl_user_op + def apply_op(self, op, other, flip=False, *, loc=None, ip=None) -> "TensorSSA": + """ + Apply a binary operation to this tensor and another operand. + + This is a public interface to the internal _apply_op method, providing + a stable API for external users who need to apply custom operations. + + Args: + op: The operation function (e.g., operator.add, operator.mul, etc.) + other: The other operand (TensorSSA, ArithValue, or scalar) + flip: Whether to flip the operands (for right-hand operations) + loc: MLIR location (optional) + ip: MLIR insertion point (optional) + + Returns: + TensorSSA: The result of the operation + + Example: + >>> tensor1 = cute.Tensor(...) + >>> tensor2 = cute.Tensor(...) + >>> result = tensor1.apply_op(operator.add, tensor2) + >>> # Equivalent to: tensor1 + tensor2 + """ + return self._apply_op(op, other, flip=flip, loc=loc, ip=ip) + @dsl_user_op def broadcast_to(self, target_shape: Shape, *, loc=None, ip=None) -> "TensorSSA": """ @@ -1147,20 +1252,21 @@ class TensorSSA(cutlass_arith.ArithValue): transform_leaf(_check_broadcast, shape, target_shape) + # convert TensorSSA col-major vec to row-m to be compatible with mlir vector ops + row_major_vec = _col2row(self, shape=shape, loc=loc, ip=ip) # reshape to flatten N-D vector flat_shp = flatten_to_tuple(shape) temp_ty = ir.VectorType.get(list(flat_shp), self.dtype.mlir_type) - temp_vect = vector.shape_cast(temp_ty, self, loc=loc, ip=ip) + temp_vect = vector.shape_cast(temp_ty, row_major_vec, loc=loc, ip=ip) # broadcast to result N-D vector flat_tgt_shp = flatten_to_tuple(target_shape) temp_tgt_ty = ir.VectorType.get(list(flat_tgt_shp), self.dtype.mlir_type) temp_tgt_vect = vector.broadcast(temp_tgt_ty, temp_vect, loc=loc, ip=ip) - res_1d_ty = ir.VectorType.get([size(target_shape)], self.dtype.mlir_type) # type: ignore - res_1d_vect = vector.shape_cast(res_1d_ty, temp_tgt_vect, loc=loc, ip=ip) - - return TensorSSA(res_1d_vect, target_shape, self.dtype) + return self._build_result( + temp_tgt_vect, target_shape, row_major=True, loc=loc, ip=ip + ) @dsl_user_op def __pow__(self, other, *, loc=None, ip=None) -> "TensorSSA": @@ -1499,15 +1605,43 @@ class TensorSSA(cutlass_arith.ArithValue): assert isinstance(flat_crd, tuple) and is_static(flat_crd) return flat_shp, flat_crd - def _build_result(self, res_vect, res_shp, *, loc=None, ip=None): + def _build_result(self, res_vect, res_shp, *, row_major=False, loc=None, ip=None): if isinstance(res_shp, ir.Value): raise ValueError(f"Expected static shape, but got {self._shape}") # cast back to 1D vector res_1d_ty = ir.VectorType.get([size(res_shp)], self.type.element_type) res_1d_vect = vector.shape_cast(res_1d_ty, res_vect, loc=loc, ip=ip) + + if row_major: + res_1d_vect = _row2col(res_1d_vect, shape=res_shp, loc=loc, ip=ip) + return TensorSSA(res_1d_vect, res_shp, self.dtype) + @dsl_user_op + def reshape(self, shape: Shape, *, loc=None, ip=None) -> "TensorSSA": + """Reshape the tensor to a new shape. + + :param shape: The new shape to reshape to. + :type shape: Shape + :return: A new tensor with the same elements but with the new shape. + :rtype: TensorSSA + :raises NotImplementedError: If dynamic size is not supported + :raises ValueError: If the new shape is not compatible with the current shape + """ + + cur_size = size(self.shape, loc=loc, ip=ip) + shp_size = size(shape, loc=loc, ip=ip) + + if type(shp_size) is not int: + raise NotImplementedError(f"dynamic shape is not supported: {shape}") + if cur_size != shp_size: + raise ValueError( + f"expected reshaped size to be the same: {self.shape} -> {shape}" + ) + + return TensorSSA(self, shape, self.dtype) + @dsl_user_op def __getitem__( self, crd: Coord, *, loc=None, ip=None @@ -1562,9 +1696,11 @@ class TensorSSA(cutlass_arith.ArithValue): flat_shp, flat_crd = self._flatten_shape_and_coord(crd, loc=loc, ip=ip) + # convert TensorSSA col-major vec to row-m to be compatible with mlir vector ops + row_major_vec = _col2row(self, shape=self._shape, loc=loc, ip=ip) multi_dim_ty = ir.VectorType.get(list(flat_shp), self.type.element_type) # vector -> vector - tmp_vect = vector.shape_cast(multi_dim_ty, self, loc=loc, ip=ip) + tmp_vect = vector.shape_cast(multi_dim_ty, row_major_vec, loc=loc, ip=ip) # Slice and keep dims matching `_` or None res_shp = slice_(self._shape, crd, loc=loc, ip=ip) @@ -1593,7 +1729,7 @@ class TensorSSA(cutlass_arith.ArithValue): # Slice and keep dims matching `_` or None res_shp = slice_(self._shape, crd, loc=loc, ip=ip) - return self._build_result(res_vect, res_shp, loc=loc, ip=ip) + return self._build_result(res_vect, res_shp, row_major=True, loc=loc, ip=ip) @dsl_user_op def to(self, dtype: Type[Numeric], *, loc=None, ip=None): @@ -1634,11 +1770,15 @@ class TensorSSA(cutlass_arith.ArithValue): src, dtype.signed, dtype.mlir_type, loc=loc, ip=ip ) elif issubclass(src_dtype, Integer) and dtype.is_float: - # fast conversion path for supported combinations + # check if there is a fast conversion path for given data types and arch + fast_cvt_func = None if src_dtype in (Int8, Uint8) and dtype == BFloat16: - res_vect = cvt_i8_bf16_intrinsic(src, size(self.shape), loc=loc, ip=ip) + fast_cvt_func = cvt_i8_bf16_intrinsic elif src_dtype == Int4 and dtype == BFloat16: - res_vect = cvt_i4_bf16_intrinsic(src, size(self.shape), loc=loc, ip=ip) + fast_cvt_func = cvt_i4_bf16_intrinsic + arch = CuTeDSL._get_dsl().get_arch_enum() + if fast_cvt_func is not None and arch in fast_cvt_func.supported_archs: + res_vect = fast_cvt_func(src, size(self.shape), loc=loc, ip=ip) else: res_vect = cutlass_arith.itofp( src, src_dtype.signed, dtype.mlir_type, loc=loc, ip=ip @@ -1748,8 +1888,10 @@ class TensorSSA(cutlass_arith.ArithValue): assert depth(flat_shp) == 1 and depth(flat_prof) == 1 assert rank(flat_shp) == rank(flat_prof) + # convert TensorSSA col-major vec to row-m to be compatible with mlir vector ops + row_major_vec = _col2row(self, shape=self._shape, loc=loc, ip=ip) temp_ty = ir.VectorType.get(list(flat_shp), elem_type.mlir_type) - temp_vect = vector.shape_cast(temp_ty, self, loc=loc, ip=ip) + temp_vect = vector.shape_cast(temp_ty, row_major_vec, loc=loc, ip=ip) red_dims = [i for i, x in enumerate(flat_prof) if x is not None] @@ -1763,7 +1905,7 @@ class TensorSSA(cutlass_arith.ArithValue): # Slice and keep dims matching `_` or None res_shp = slice_(self.shape, reduction_profile, loc=loc, ip=ip) - return self._build_result(res_vect, res_shp, loc=loc, ip=ip) + return self._build_result(res_vect, res_shp, row_major=True, loc=loc, ip=ip) @dsl_user_op diff --git a/python/CuTeDSL/cutlass/cute/testing.py b/python/CuTeDSL/cutlass/cute/testing.py index ebb9824c..6f3d7e17 100644 --- a/python/CuTeDSL/cutlass/cute/testing.py +++ b/python/CuTeDSL/cutlass/cute/testing.py @@ -1057,7 +1057,7 @@ def tune( except NotImplementedError as e: logger.info(f" Encountered unimplemented error, abort execution: {e}") raise e - except (ValueError, TypeError) as e: + except (ValueError, TypeError, CantImplementError) as e: logger.info(f" Configuration parameter skipping: {e}") continue except Exception as e: @@ -1073,3 +1073,17 @@ def tune( logger.info(f"Best configuration: {best_config}, execution time: {min_time} us") logger.info(f"Total tuning time: {tuning_time} s") return best_config + + +class CantImplementError(Exception): + """Exception raised when a function is not implemented.""" + + def __init__(self, message=None): + self.message = message or "The current config is invalid/unsupported" + super().__init__(self.message) + + def __str__(self): + return self.message + + def __repr__(self): + return self.message diff --git a/python/CuTeDSL/cutlass/cute/typing.py b/python/CuTeDSL/cutlass/cute/typing.py index 3048f550..d61c3c90 100644 --- a/python/CuTeDSL/cutlass/cute/typing.py +++ b/python/CuTeDSL/cutlass/cute/typing.py @@ -10,17 +10,79 @@ # is strictly prohibited. from abc import ABC, abstractmethod -from typing import ForwardRef, Tuple, Union, Any, Type, List, Optional +import ctypes +from typing import ForwardRef, Tuple, Union, Any, Type, List, Optional, Literal +from functools import lru_cache from cutlass.base_dsl.typing import * from cutlass._mlir import ir from cutlass._mlir.dialects.cute import AddressSpace, ConstrainedIntType +from cutlass.base_dsl.typing import JitArgument Int = Union[int, Integer] +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 + + @property + def width(self): + return self._width + + @property + def divisibility(self): + return self._divisibility + + def __str__(self) -> str: + return f"?{{i{self._width} div={self._divisibility}}}" + + def __repr__(self) -> str: + return self.__str__() + + def __eq__(self, other) -> bool: + if not isinstance(other, SymInt): + return False + return all( + [self._width == other._width, self._divisibility == other._divisibility] + ) + + def __c_pointers__(self): + return [ctypes.c_void_p(0).value] + + def __get_mlir_types__(self) -> List[ir.Type]: + res_ty = ir.Type.parse( + f'!cute.int_tuple<"?{{i{self.width} div={self.divisibility}}}">' + ) + return [res_ty] + + def __new_from_mlir_values__(self, values) -> "SymInt": + from .core import IntValue + + if self.width == 32: + return Int32(IntValue(values[0])) + elif self.width == 64: + return Int64(IntValue(values[0])) + 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) + + +def sym_int32(divisibility=1) -> SymInt: + return sym_int(32, divisibility=divisibility) + + +def sym_int64(divisibility=1) -> SymInt: + return sym_int(64, divisibility=divisibility) + + ScaledBasis = ForwardRef("ScaledBasis") IntTuple = Union[Int, Tuple["IntTuple", ...]] @@ -288,7 +350,10 @@ def is_int_tuple(a) -> bool: __all__ = [ - "Coord", + "SymInt", + "sym_int", + "sym_int32", + "sym_int64", "Numeric", "Integer", "Boolean", @@ -317,13 +382,14 @@ __all__ = [ "Float6E3M2FN", "IntTuple", "Layout", - "Pointer", + "Coord", "Shape", "Stride", - "Tensor", "Tile", "Tiler", "XTuple", "is_integer", "is_int_tuple", + "Pointer", + "Tensor", ] diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py b/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py new file mode 100644 index 00000000..08321ee5 --- /dev/null +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py @@ -0,0 +1,333 @@ +# 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 +# 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. + +""" +This module provides jit executor related classes for CUTLASS. +""" +import ctypes +import functools +import weakref +import threading + +import cuda.bindings.runtime as cuda_runtime +import cuda.bindings.driver as cuda_driver + +# Local modules imports +from ..base_dsl.jit_executor import JitExecutor, ExecutionArgs, JitFunctionArtifacts +from ..base_dsl.utils.logger import log +from ..base_dsl.common import DSLCudaRuntimeError, DSLRuntimeError +from ..base_dsl.typing import Int32 + +class CudaDialectJitModule: + """Holds the execution engine and cuda libraries.""" + + def __init__( + self, + engine, + capi_func, + args_spec: ExecutionArgs, + cuda_library: list["cuda_runtime.cudaLibrary_t"], + ): + self.engine = engine + self.capi_func = capi_func + self.args_spec = args_spec + self.cuda_library = cuda_library + self._unloaded = False + + def is_unloaded(self): + return self._unloaded + + def unload(self): + try: + for library in self.cuda_library: + cuda_runtime.cudaLibraryUnload(library) + self.cuda_library.clear() + finally: + self._unloaded = True + + def __del__(self): + self.unload() + + +class CudaDialectJitCompiledFunction: + """Holds a compiled function and its module.""" + + def __init__( + self, + ir_module, + engine, + capi_func, + args_spec, + function_name, + kernel_info, + jit_time_profiling, + jit_function_artifacts, + prefix=None, + load_from_binary=False, + ): + self.ir_module = ir_module + self.engine = engine + self.capi_func = capi_func + self.function_name = function_name + self.kernel_info = kernel_info + if args_spec is not None: + self.args_spec = ExecutionArgs(args_spec, self.function_name) + self.jit_time_profiling = jit_time_profiling + self.prefix = prefix + assert ( + isinstance(jit_function_artifacts, JitFunctionArtifacts) + or jit_function_artifacts is None + ) + self.artifacts = jit_function_artifacts + self.load_from_binary = load_from_binary + + # Set cuda result return type. + # When execution engine/capi function is None, do not set the return type. + if self.capi_func: + self.capi_func.restype = ctypes.c_int32 + if self.args_spec: + self.args_spec.args_spec.annotations["return"] = Int32 + + # This runtime state is stored here so that we can preserve the module + # in the compiler cache. Callers can extend the lifetime of the module + # by creating and retaining the executor. + self.jit_module = None + self._executor_lock = threading.RLock() + self._default_executor = None + + @property + def __ptx__(self): + """Returns the PTX code of the JIT-compiled function.""" + return self.artifacts.PTX if self.artifacts is not None else None + + @property + def __cubin__(self): + """Returns the CUBIN data of the JIT-compiled function.""" + return self.artifacts.CUBIN if self.artifacts is not None else None + + @property + def __mlir__(self): + """Returns the MLIR code of the JIT-compiled function.""" + return self.artifacts.MLIR if self.artifacts is not None else None + + @functools.cached_property + def num_devices(self): + """Returns the number of CUDA devices available.""" + dev_err, devs = cuda_runtime.cudaGetDeviceCount() + if dev_err != cuda_runtime.cudaError_t.cudaSuccess: + raise DSLCudaRuntimeError(dev_err, cuda_runtime.cudaGetErrorName(dev_err)) + return devs + + def _deserializer(self): + """Load the cuda library from the binary execution engine. + @return: The list of cuda kernels. + """ + library = ctypes.c_void_p() + pointer_to_library = ctypes.pointer(library) + pointer_to_pointer_to_library = ctypes.pointer(pointer_to_library) + err = ctypes.c_int32(0) + pointer_to_err = ctypes.pointer(err) + + # cuda init takes in a pointer to a cudaLibrary_t and returns + # a i32 cudaError_t. It initialized (lazy loads) our cudaLibrary_t + cuda_init = self.engine.lookup(f"_mlir_{self.prefix}_cuda_init") + if cuda_init is None: + raise DSLRuntimeError("cuda_init not found") + cuda_init = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(cuda_init) + # cuda load takes in a pointer to a cudaLibrary_t and returns + # a i32 cudaError_t. It loads the functions from the cuda library, + # sets function attributes, and returns an error if encountered. + cuda_load = self.engine.lookup(f"_mlir_{self.prefix}_cuda_load") + if cuda_load is None: + raise DSLRuntimeError("cuda_load not found") + cuda_load = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(cuda_load) + + cuda_init_args = [pointer_to_pointer_to_library, pointer_to_err] + packed_args = (ctypes.c_void_p * len(cuda_init_args))() + for i in range(len(cuda_init_args)): + packed_args[i] = ctypes.cast(cuda_init_args[i], ctypes.c_void_p) + cuda_init(packed_args) + + if err.value != 0: + error_code = err.value + error_name = cuda_runtime.cudaGetErrorName( + cuda_runtime.cudaError_t(error_code) + ) + raise DSLCudaRuntimeError(error_code, error_name) + + cuda_load_args = [pointer_to_library, pointer_to_err] + packed_args = (ctypes.c_void_p * len(cuda_load_args))() + for i in range(len(cuda_load_args)): + packed_args[i] = ctypes.cast(cuda_load_args[i], ctypes.c_void_p) + cuda_load(packed_args) + + if err.value != 0: + error_code = err.value + error_name = cuda_runtime.cudaGetErrorName( + cuda_runtime.cudaError_t(error_code) + ) + raise DSLCudaRuntimeError(error_code, error_name) + + return [cuda_runtime.cudaLibrary_t(library.value)] + + def _get_cuda_init_and_load(self): + """Returns the cuda init and load functions from the engine.""" + # cuda init takes in a pointer to a cudaLibrary_t and returns + # a i32 cudaError_t. It initialized (lazy loads) our cudaLibrary_t + + cuda_init = None + + # cuda load for device takes in a pointer to a cudaLibrary_t and a device index and returns + # a i32 cudaError_t. It resolves each kernel from the cuda library, + # and applies device scoped attributes for the given device id, and returns an error if encountered. + cuda_load_to_device = None + + # When load_from_binary is true, the symbols are prefixed by _mlir__ and are looked + # up from the JIT engine. Otherwise we look for the unprefixed forms. Looking up cuda_init + # and cuda_load_to_device from the engine which are defined in CudaToLLVM. + if self.load_from_binary: + if self.prefix is None: + raise DSLRuntimeError("prefix is required to be set for binary loading") + cuda_init = self.engine.lookup(f"_mlir_{self.prefix}_cuda_init") + if cuda_init is None: + raise DSLRuntimeError(f"cuda_init not found for prefix {self.prefix}") + cuda_load_to_device = self.engine.lookup( + f"_mlir_{self.prefix}_cuda_load_to_device" + ) + if cuda_load_to_device is None: + raise DSLRuntimeError( + f"cuda_load_to_device not found for prefix {self.prefix}" + ) + else: + cuda_init = self.engine.raw_lookup("cuda_init") + if cuda_init is None: + raise DSLRuntimeError("cuda_init not found") + cuda_load_to_device = self.engine.raw_lookup("cuda_load_to_device") + if cuda_load_to_device is None: + raise DSLRuntimeError("cuda_load_to_device not found") + cuda_init = ctypes.CFUNCTYPE(ctypes.c_int32, ctypes.c_void_p)(cuda_init) + cuda_load_to_device = ctypes.CFUNCTYPE(ctypes.c_int32, ctypes.c_void_p)( + cuda_load_to_device + ) + + return cuda_init, cuda_load_to_device + + def _load_cuda_library(self): + """Loads the CUDA library from the engine.""" + + cuda_init, cuda_load_to_device = self._get_cuda_init_and_load() + + library = ctypes.c_void_p() + pointer_to_library = ctypes.pointer(library) + pointer_to_pointer_to_library = ctypes.pointer(pointer_to_library) + err = ctypes.c_int32(0) + pointer_to_err = ctypes.pointer(err) + + cuda_init_args = [pointer_to_pointer_to_library, pointer_to_err] + packed_args = (ctypes.c_void_p * len(cuda_init_args))() + for i in range(len(cuda_init_args)): + packed_args[i] = ctypes.cast(cuda_init_args[i], ctypes.c_void_p) + cuda_init(packed_args) + + if err.value != 0: + error_code = err.value + error_name = cuda_runtime.cudaGetErrorName( + cuda_runtime.cudaError_t(error_code) + ) + raise DSLCudaRuntimeError(error_code, error_name) + + device_id = ctypes.c_int32(0) + pointer_to_device_id = ctypes.pointer(device_id) + + cuda_load_args = [pointer_to_library, pointer_to_device_id, pointer_to_err] + packed_args = (ctypes.c_void_p * len(cuda_load_args))() + for i, arg in enumerate(cuda_load_args): + packed_args[i] = ctypes.cast(arg, ctypes.c_void_p) + + for dev in range(self.num_devices): + device_id.value = dev + cuda_load_to_device(packed_args) + if err.value != 0: + raise DSLCudaRuntimeError( + err.value, + cuda_runtime.cudaGetErrorName(cuda_runtime.cudaError_t(err.value)), + ) + + if err.value != 0: + error_code = err.value + error_name = cuda_runtime.cudaGetErrorName( + cuda_runtime.cudaError_t(error_code) + ) + raise DSLCudaRuntimeError(error_code, error_name) + + return [cuda_runtime.cudaLibrary_t(library.value)] + + def to(self, device=None) -> JitExecutor: + """Returns an executable function bound to the given device. + + For multi-device execution this method can be called for each device where + the kernel will run. + + Since CudaJitCompiledFunction uses CUDA libraries, which are context free, + binding to a device is not necessary and the device is ignored. Device is + kept in for compatibility with the JitCompiledFunction. + + :param device: Specifies the device for the executor. If None the current device is used. + :type device: Optional[Union[int, CUdevice]] + :return: A callable executor function. + :rtype: JitExecutor + """ + with self._executor_lock: + # We need to ensure that the modules are loaded if not already + if self.jit_module is None or self.jit_module.is_unloaded(): + cuda_library = self._load_cuda_library() + self.jit_module = CudaDialectJitModule( + self.engine, + self.capi_func, + self.args_spec, + cuda_library, + ) + + return JitExecutor(self.jit_module, None, self.jit_time_profiling) + + def generate_execution_args(self, *args, **kwargs): + return self.args_spec.generate_execution_args(args, kwargs) + + def set_dynamic_args(self, dynamic_args, dynamic_kwargs): + """Sets the dynamic argument information required for export to c code generation.""" + self.dynamic_args = dynamic_args + self.dynamic_kwargs = dynamic_kwargs + + def __call__(self, *args, **kwargs): + """Executes the jit-compiled function under the currently active CUDA context. + + Calling this method multiple devices is not allowed and will result in unexpected + CUDA errors. If you need to call the kernel on multiple devices use `to` + to return a per-device function. + """ + exe_args, adapted_args = self.generate_execution_args(*args, **kwargs) + return self.run_compiled_program(exe_args) + + def run_compiled_program(self, exe_args): + """Executes the jit-compiled function under the currently active CUDA context. + + Calling this method multiple devices is not allowed and will result in unexpected + CUDA errors. If you need to call the kernel on multiple devices use `to` + to return a per-device function. + """ + with self._executor_lock: + if self._default_executor is None: + log().debug("Creating default executor.") + # We use a weak reference here so that this instance does not keep this + # object alive as it hold a reference to self. + proxy_self = weakref.proxy(self) + self._default_executor = proxy_self.to(None) + return self._default_executor.run_compiled_program(exe_args) diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cuda_stream_adapter.py b/python/CuTeDSL/cutlass/cutlass_dsl/cuda_stream_adapter.py new file mode 100644 index 00000000..05d3d6b4 --- /dev/null +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cuda_stream_adapter.py @@ -0,0 +1,44 @@ +# 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 +# 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. + +""" +This module provides CUDA Python helper functions +""" + +import cuda.bindings.driver as cuda_driver + +# MLIR imports +from .._mlir import ir +from .._mlir.dialects import cuda + +# Local module imports +from ..base_dsl.runtime.jit_arg_adapters import JitArgAdapterRegistry + + +@JitArgAdapterRegistry.register_jit_arg_adapter(cuda_driver.CUstream) +class CudaDialectStreamAdapter: + """ + Convert a CUDA stream to a stream representation for JIT arg generation. + """ + + def __init__(self, arg): + self._arg = arg + self._c_pointer = self._arg.getPtr() + + def __new_from_mlir_values__(self, values): + assert len(values) == 1 + return values[0] + + def __c_pointers__(self): + return [self._c_pointer] + + def __get_mlir_types__(self): + return [cuda.StreamType.get()] diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py index 19684d9b..cf8162e6 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py @@ -43,11 +43,21 @@ from ..base_dsl.typing import * from ..base_dsl.typing import DynamicExpression, get_mlir_types from ..base_dsl.runtime.jit_arg_adapters import is_arg_spec_constexpr from ..base_dsl.runtime import cuda as cuda_helpers +from .cuda_stream_adapter import CudaDialectStreamAdapter +from .cuda_jit_executor import CudaDialectJitCompiledFunction # MLIR Imports from cutlass._mlir import ir, execution_engine, passmanager -from cutlass._mlir.dialects import arith, func, gpu, scf, cute, gpu as cutlass_gpu +from cutlass._mlir.dialects import ( + arith, + func, + gpu, + scf, + cute, + gpu as cutlass_gpu, + cuda as cuda_dialect, +) from cutlass._mlir.dialects._ods_common import ( get_op_result_or_op_results as _get_op_result_or_op_results, ) @@ -220,6 +230,10 @@ class CutlassBaseDSL(BaseDSL): preprocess=preprocess, ) self._smem_usage_tracker: tuple = None + # extra function to convert cute arguments to tvm ffi spec params + # this needs to be reverse registered because the arg convention + # depends on the runtime type of the DSL arguments + self._tvm_ffi_args_spec_converter = None # this method is not useful for cutlass_dsl, so we only provide a dummy implementation. def _is_tensor_descriptor(self, maybe_tensor_descriptor) -> bool: @@ -241,7 +255,7 @@ class CutlassBaseDSL(BaseDSL): def _get_pipeline(self, pipeline): pipeline = super()._get_pipeline(pipeline) - if pipeline == None: + if pipeline is None: # cubin format is required to be cubin as we launch cuda module at python level. return ( "builtin.module(cute-to-nvvm{cubin-format=bin " @@ -253,7 +267,10 @@ class CutlassBaseDSL(BaseDSL): def preprocess_pipeline(self, pipeline, arch) -> str: pipeline = super().preprocess_pipeline(pipeline, arch) - pipeline = pipeline.rstrip(")") + ",external-kernel-for-gpu-launch)" + pipeline = ( + pipeline.rstrip("})") + + " enable-cuda-dialect=true cuda-dialect-external-module=true})" + ) return pipeline def _enter_gpu_module(self): @@ -265,13 +282,35 @@ class CutlassBaseDSL(BaseDSL): ) ret = {} - # generate launch bound attr from LaunchConfig - max_threads = ", ".join(map(str, config.block)) - ret["nvvm.reqntid"] = ir.Attribute.parse(f"array") - # min_blocks_per_mp is optional for kernel - min_blocks = config.min_blocks_per_mp - if min_blocks > 0: - ret["nvvm.minctasm"] = ir.Attribute.parse(f"{min_blocks} : i32") + if config.has_max_number_threads(): + block_str = ", ".join(map(str, config.block)) + has_dynamic = any(is_dynamic_expression(dim) for dim in config.block) + if not has_dynamic: + # Case 1. Auto-generate reqntid from block size + ret["nvvm.reqntid"] = ir.Attribute.parse(f"array") + else: + # Case 2. Don't auto-generate reqntid from block size + self.print_warning_once( + f"Dynamic variable in block size {block_str}, cannot auto-generate `nvvm.reqntid`" + ) + else: + # Case 3. Use user-defined value for maxntid + max_ntid_str = ", ".join(map(str, config.max_number_threads)) + for value in config.max_number_threads: + if is_dynamic_expression(value): + raise DSLRuntimeError( + f"User defined max number threads `{max_ntid_str}` contains dynamic expression, cannot generate `nvvm.maxntid`", + suggestion="Consider using `Constexpr` annotation or Python constant", + ) + # Use user-defined maxntid + ret["nvvm.maxntid"] = ir.Attribute.parse(f"array") + + # Add optional minimum blocks per multiprocessor + if config.min_blocks_per_mp > 0: + ret["nvvm.minctasm"] = ir.Attribute.parse( + f"{config.min_blocks_per_mp} : i32" + ) + return ret @lru_cache(maxsize=1) @@ -313,6 +352,116 @@ class CutlassBaseDSL(BaseDSL): return version_hash + def get_return_types(self) -> List[ir.Type]: + """ + Get the return types of the function. + With cuda dialect, the return type is i32 to indicate cuda result. + """ + i32_ty = ir.IntegerType.get_signless(32) + return [i32_ty] + + def generate_default_return_values(self, ip=None) -> List[ir.Value]: + """ + Generate the default return values of the function. + With cuda dialect, the default return value is 0 to indicate success. + """ + if ip is None: + raise DSLRuntimeError("ip is required to generate default return values") + return_types = self.get_return_types() + if len(return_types) != 1: + raise DSLRuntimeError( + f"Expected exactly one return type, but got {len(return_types)}" + ) + return_type = return_types[0] + with ip: + zero = arith.ConstantOp(return_type, 0).result + return [zero] + + def compile_and_cache( + self, + module, + module_hash, + function_name, + pipeline, + args_spec, + no_cache, + *, + dynamic_args=None, + dynamic_kwargs=None, + original_function_name=None, + ): + """ + Compile the module and cache the result. + + :param module: The MLIR module to compile. + :param module_hash: The hash of the MLIR module. + :param function_name: The name of the function to compile. + :param pipeline: The pipeline to use for compilation. + :param args_spec: The args spec to use for compilation. + :param no_cache: Whether to cache the result. + :param dynamic_args: The dynamic arguments to use for compilation. + :param dynamic_kwargs: The dynamic keyword arguments to use for compilation. + :param original_function_name: The name of the original function without mangling. + :return: The compiled function. + """ + if self.compile_options.enable_tvm_ffi: + # TVM FFI post compile logic + # attach extra ABI function to the MLIR module + from .tvm_ffi_provider import ( + TVMFFIJitCompiledFunction, + TVMFFICuteCallProvider, + ) + from cutlass.base_dsl.tvm_ffi_builder import attach_ffi_func + + assert self._tvm_ffi_args_spec_converter is not None + tvm_ffi_spec_params = self._tvm_ffi_args_spec_converter( + function_name, args_spec, dynamic_args, dynamic_kwargs + ) + tvm_ffi_provider = TVMFFICuteCallProvider(function_name) + + # ensure we run the postprocessor hook after the compiler has run its passes + def post_compile_hook(module: ir.Module): + with module.context, module.operation.location: + # attach the tvm ffi function to the mlir module + attach_ffi_func( + module, + function_name, + tvm_ffi_spec_params, + tvm_ffi_provider, + fn_display_name=original_function_name, + ) + module.operation.verify() + + # ensure the compiler can run post-compile hook after its passes + # the context will restore the previous post-compile hook after it exits + with compiler.PostCompileHookContext( + self.compiler_provider, post_compile_hook + ): + return super().compile_and_cache( + module, + module_hash, + function_name, + pipeline, + args_spec, + no_cache, + TVMFFIJitCompiledFunction, + dynamic_args=dynamic_args, + dynamic_kwargs=dynamic_kwargs, + ) + + return super().compile_and_cache( + module, + module_hash, + function_name, + pipeline, + args_spec, + no_cache, + CudaDialectJitCompiledFunction, + dynamic_args=dynamic_args, + dynamic_kwargs=dynamic_kwargs, + original_function_name=original_function_name, + ) + @staticmethod def track_smem_allocator(allocator, callback): """ @@ -343,6 +492,103 @@ class CutlassBaseDSL(BaseDSL): allocator, callback = self._smem_usage_tracker return callback(allocator) + @staticmethod + def cuda_launch_func( + stream: Union[list, tuple], + kernel, + grid_size_x, + grid_size_y, + grid_size_z, + block_size_x, + block_size_y, + block_size_z, + kernel_operands, + *, + cluster_size_x=None, + cluster_size_y=None, + cluster_size_z=None, + dynamic_shared_memory_size=None, + use_pdl=False, + loc=None, + ip=None, + ): + + # Max number of attributes in the launch config is set to 2 for PDL and cluster size + max_num_attributes = 2 + launch_config_type = cuda_dialect.LaunchConfigType.get(max_num_attributes) + + if len(stream) == 0: + stream = cuda_dialect.CastOp( + src=Int64(0).ir_value(), + dst=cuda_dialect.StreamType.get(), + loc=loc, + ip=ip, + ) + else: + stream = stream[0] + + dynamic_shared_memory_size = Int64(dynamic_shared_memory_size).ir_value() + block_size_x = Int32(block_size_x).ir_value() + block_size_y = Int32(block_size_y).ir_value() + block_size_z = Int32(block_size_z).ir_value() + grid_size_x = Int32(grid_size_x).ir_value() + grid_size_y = Int32(grid_size_y).ir_value() + grid_size_z = Int32(grid_size_z).ir_value() + + cfg = cuda_dialect.launch_cfg_create( + # Launch config type + launch_config_type, + # Max num of attributes the launch config can hold + # set to 2 for PDL and cluster size + ir.IntegerAttr.get(ir.IntegerType.get_signless(32), max_num_attributes), + block_size_x, + block_size_y, + block_size_z, + dynamic_shared_memory_size, + grid_size_x, + grid_size_y, + grid_size_z, + stream, + loc=loc, + ip=ip, + ) + + if use_pdl: + cuda_dialect.launch_cfg_programmatic_stream_serialization_allowed( + cfg, Int32(use_pdl).ir_value(), loc=loc, ip=ip + ) + + if cluster_size_x is not None: + if cluster_size_y is None: + cluster_size_y = 1 + if cluster_size_z is None: + cluster_size_z = 1 + cluster_size_x = Int32(cluster_size_x).ir_value(loc=loc, ip=ip) + cluster_size_y = Int32(cluster_size_y).ir_value(loc=loc, ip=ip) + cluster_size_z = Int32(cluster_size_z).ir_value(loc=loc, ip=ip) + cuda_dialect.launch_cfg_cluster_dim( + cfg, cluster_size_x, cluster_size_y, cluster_size_z, loc=loc, ip=ip + ) + + op = cuda_dialect.launch_ex( + cuda_dialect.ResultType.get(), + kernel, + cfg, + kernel_operands, + loc=loc, + ip=ip, + ) + + res_i32 = cuda_dialect.CastOp( + src=op, + dst=cuda_dialect.IntegerType.get_signless(32), + loc=loc, + ip=ip, + ) + + cuda_dialect.return_if_error([res_i32], loc=loc, ip=ip) + return None + @staticmethod def gpu_launch_func( async_token, @@ -396,20 +642,32 @@ class CutlassBaseDSL(BaseDSL): def generate_func_op(self, arg_types, arg_attrs, kernel_name, loc=None): super().generate_func_op(arg_types, arg_attrs, kernel_name) - self.func_op = func.FuncOp( + self.func_op = cuda_dialect.KernelOp( kernel_name, ir.FunctionType.get(arg_types, []), loc=loc ) + self.arg_types = arg_types + self.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(), + } + ) if arg_attrs is not None: log().debug(arg_attrs) self.func_op.arg_attrs = arg_attrs return self.func_op - def generate_func_ret_op(self): - return func.ReturnOp([]) + def generate_func_ret_op(self, loc=None, ip=None): + return cuda_dialect.ReturnOp([], loc=loc, ip=ip) def get_func_body_start(self): assert self.func_op is not None, "Invalid func_op is not expected!" - return self.func_op.add_entry_block() + arg_locs = [self.func_op.operation.location for _ in self.arg_types] + return self.func_op.add_entry_block(arg_locs=arg_locs) def generate_launch_op(self, *args, **kwargs): # Extract args and do validation @@ -430,12 +688,6 @@ class CutlassBaseDSL(BaseDSL): cfg = requiredArgs.config - # Apply to grid, block, and cluster if present - cfg.grid = [to_index(size) for size in cfg.grid] - cfg.block = [to_index(size) for size in cfg.block] - if cfg.has_cluster: - cfg.cluster = [to_index(size) for size in cfg.cluster] - smem_usage = self.dsl._get_smem_usage() if any(not isinstance(x, int) for x in [cfg.smem, smem_usage]): pass # cannot compare dynamic value inside kernel to launch op in py @@ -449,12 +701,11 @@ class CutlassBaseDSL(BaseDSL): ) cfg.smem = const(cfg.smem) + async_deps = cfg.async_deps if not isinstance(cfg.async_deps, (list, tuple)): - cfg.async_deps = [cfg.async_deps] - is_async = len(cfg.async_deps) > 0 - token = CutlassBaseDSL.gpu_launch_func( - gpu.AsyncTokenType.get() if is_async else None, - cfg.async_deps, + async_deps = [cfg.async_deps] + CutlassBaseDSL.cuda_launch_func( + async_deps, kernelSym, *cfg.grid, *cfg.block, @@ -469,7 +720,7 @@ class CutlassBaseDSL(BaseDSL): use_pdl=cfg.use_pdl, loc=loc, ) - return token if is_async else None + return None return KernelLauncher( self, @@ -526,6 +777,10 @@ class CutlassBaseDSL(BaseDSL): # skip generic types such as List[int], Tuple[int, int], etc. for performance consideration? pass + elif getattr(arg_annotation, "__name__", None) == "CUstream" and arg.__class__.__name__ == "_FakeStream": + # allow FakeStream to be passed as CUDA stream + pass + elif isinstance(arg_annotation, type): # Handle simple type annotations if not isinstance(arg, arg_annotation) and arg is not None: @@ -564,7 +819,23 @@ class CutlassBaseDSL(BaseDSL): # Handle dynamic types jit_arg_type.extend([v.type for v in dyn_vals]) jit_arg_attr.extend([default_attr] * len(dyn_vals)) - jit_exec_arg.extend(get_c_pointers(arg) if is_host else dyn_vals) + if is_host and self.envar.enable_tvm_ffi: + import tvm_ffi + + jit_exec_arg.extend( + [ + tvm_ffi.Shape( + [ + v.value if isinstance(v, Numeric) else v + for v in arg + ] + ) + ] + ) + else: + jit_exec_arg.extend( + get_c_pointers(arg) if is_host else dyn_vals + ) else: jit_exec_arg = jit_arg_type = jit_arg_attr = None elif not hasattr(arg, "__extract_mlir_values__") and not hasattr( @@ -879,7 +1150,6 @@ def to_index(value): return res - def _validate_iter_args_structure(iter_args, ir_values): """ Validates that iter_args structure contains the same number of atomic values @@ -938,6 +1208,15 @@ def _minmax(op, *args, loc=None, ip=None): emitter = getattr(cutlass_arith, f"_{op.__name__}") if not (is_dynamic_expression(res) or is_dynamic_expression(x)): res = emitter(op(res), op(x)) + elif ( + hasattr(res, "type") + and hasattr(x, "type") + and isinstance(res.type, T.VectorType) + and isinstance(x.type, T.VectorType) + ): + lhs = op(res, loc=loc, ip=ip) + rhs = op(x, loc=loc, ip=ip) + return lhs.apply_op(emitter, rhs, loc=loc, ip=ip) else: lhs = as_numeric(op(res, loc=loc, ip=ip)) rhs = as_numeric(op(x, loc=loc, ip=ip)) diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py b/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py new file mode 100644 index 00000000..0a430543 --- /dev/null +++ b/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py @@ -0,0 +1,348 @@ +# 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 +# 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.base_dsl.tvm_ffi_builder import ( + DynamicParamPackCallProvider, + CallContext, + rename_tvm_ffi_function, + spec, +) +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass._mlir._mlir_libs import _execution_engine_extra +from cutlass.cutlass_dsl.cuda_jit_executor import CudaDialectJitCompiledFunction +from cutlass.base_dsl.common import DSLRuntimeError +import tvm_ffi + + +class TVMFFICuteCallProvider(DynamicParamPackCallProvider): + """Cute call provider that uses cute call convention.""" + + def __init__(self, target_func: str): + super().__init__(target_func, struct_call=True) + self.cuda_global_state_symbol = f"__{target_func}_cuda_state" + + def get_callee_struct_for_param_tensor( + self, + param: spec.Tensor, + current_block: ir.Block, + data: ir.Value, + shape: list[ir.Value], + strides: list[ir.Value], + flatten_struct: ir.Type, + ) -> ir.Type: + """Routine used to override the tensor passing struct convention""" + with ir.InsertionPoint(current_block): + data_type = self.gpu_ptr_type + strides_type = ( + self.struct_type(fields=[x.type for x in strides]) + if len(strides) != 1 + else strides[0].type + ) + shape_type = ( + self.struct_type(fields=[x.type for x in shape]) + if len(shape) != 1 + else shape[0].type + ) + shape_stride_tuple_type = self.struct_type( + fields=[shape_type, strides_type] + ) + tensor_type = self.struct_type(fields=[data_type, shape_stride_tuple_type]) + return tensor_type + + def pack_param_shape( + self, current_block: ir.Block, context: CallContext, param: spec.Shape + ) -> tuple[tuple[ir.Type], tuple[ir.Value]]: + """Pack a shape parameter to a struct.""" + allocas: list[ir.Value] = [] + arg_types: list[ir.Type] = [] + for dim in param.shape: + if isinstance(dim, spec.Var): + allocas.append( + self.pack_values_to_alloca( + current_block, + context.entry_block, + [context.matched_var_binding[dim]], + )[1] + ) + arg_types.append(context.matched_var_binding[dim].type) + return tuple(arg_types), tuple(allocas) + + def declare_extern_funcs(self, current_block: ir.Block, context: CallContext): + """Append the error handling function to the current block.""" + with ir.InsertionPoint(context.module.body): + self.declare_extern_func( + "cuda_dialect_get_error_name", + [self.i32_type], + self.ptr_type, + ) + self.declare_extern_func( + "_cudaSetDevice", + [self.i32_type], + self.i32_type, + ) + self.declare_extern_func( + "cuda_dialect_init_library_once", + [self.ptr_type, self.ptr_type, self.ptr_type, self.ptr_type], + self.i32_type, + ) + self.declare_extern_func( + "cuda_dialect_unload_library_once", + [self.ptr_type], + self.void_type, + ) + return current_block + + def insert_lazy_init_cuda(self, current_block: ir.Block, context: CallContext): + """Insert the lazy init cuda function.""" + # create global private static that is initialized to nullptr + with ir.InsertionPoint(context.module.body): + parsed_op = ir.Operation.parse( + f"llvm.mlir.global private @{self.cuda_global_state_symbol}(0 : i64) : i64" + ) + context.module.body.append(parsed_op) + + + with ir.InsertionPoint(current_block): + cuda_global_state_ptr = self.address_of( + self.cuda_global_state_symbol, self.ptr_type + ) + cuda_init_ptr = self.address_of("cuda_init", self.ptr_type) + cuda_load_ptr = self.address_of("cuda_load", self.ptr_type) + set_error_ptr = self.address_of( + "TVMFFIErrorSetRaisedFromCStr", self.ptr_type + ) + # Call the callback function with the loaded ptr value + init_result = llvm.call( + result=self.i32_type, # function returns i32 + callee="cuda_dialect_init_library_once", + callee_operands=[ + cuda_global_state_ptr, + cuda_init_ptr, + cuda_load_ptr, + set_error_ptr, + ], + op_bundle_sizes=[], + op_bundle_operands=[], + ) + # Create blocks for conditional branching + error_block = current_block.create_after() + success_block = error_block.create_after() + # Check if initialization failed (non-zero return code) + llvm.cond_br( + self.equal(init_result, self.i32(0)), + true_dest_operands=[], + false_dest_operands=[], + true_dest=success_block, + false_dest=error_block, + ) + # Error block: return the error code + # error is already set by cuda_dialect_init_library_once + with ir.InsertionPoint(error_block): + llvm.return_(arg=self.i32(-1)) + + # Continue with success block + return success_block + + def append_unload_to_global_dtors( + self, current_block: ir.Block, context: CallContext + ) -> ir.Block: + """Append the cuda_dialect_unload_library_once function to the global destructor list.""" + unload_func_symbol = "cuda_dialect_unload_library_once" + # define a private function to call the extern function, we need this wrapper function + # since llvm.mlir.global_dtors require the dtor defined in the module + unload_func_wrapper_symbol = f"__dtor_{unload_func_symbol}" + with ir.InsertionPoint(context.module.body): + params, entry_block = self.function( + name=unload_func_wrapper_symbol, + params_type=[], + ret_type=self.void_type, + internal=True, + ) + with ir.InsertionPoint(entry_block): + llvm.call( + result=None, + callee=unload_func_symbol, + callee_operands=[ + self.address_of(self.cuda_global_state_symbol, self.ptr_type) + ], + op_bundle_sizes=[], + op_bundle_operands=[], + ) + llvm.return_() + + # find or create the global destructors + global_dtors_list: list[ir.Operation] = self.find_operations_in_module( + context.module, "llvm.mlir.global_dtors" + ) + if len(global_dtors_list) == 0: + # create the global destructors + with ir.InsertionPoint(context.module.body): + global_dtors = llvm.mlir_global_dtors( + dtors=[], + priorities=[], + ) + else: + # use the existing global destructors + global_dtors = global_dtors_list[0] + + # append the unload function to the global destructors + global_dtors.attributes["dtors"] += [ + ir.FlatSymbolRefAttr.get(unload_func_wrapper_symbol) + ] + global_dtors.attributes["priorities"] += [ + ir.IntegerAttr.get(self.i32_type, 65535) + ] # the default priority + + return current_block + + def check_cuda_error( + self, code: ir.Value, current_block: ir.Block, context: CallContext + ): + """Check if the CUDA error is raised and return the error string if so.""" + with ir.InsertionPoint(current_block): + # check if the call is successful + error_block = current_block.create_after() + success_block = error_block.create_after() + # Check if call is successful (non-zero return code) + self.cond_br( + cond=self.equal(code, self.i32(0)), + true_block=success_block, + false_block=error_block, + branch_weights=self.BRANCH_WEIGHTS_LIKELY, + ) + + # Error block: raise error and return + with ir.InsertionPoint(error_block): + error_str = llvm.call( + result=self.ptr_type, + callee="cuda_dialect_get_error_name", + callee_operands=[code], + op_bundle_sizes=[], + op_bundle_operands=[], + ) + # Raise error and return -1 + context.builder.raise_error_and_return( + error_kind="RuntimeError", + error_message_parts=["CUDA Error: ", error_str], + ) + return success_block + + def generate_llvm_call( + self, + current_block: ir.Block, + call_operands: list[ir.Value], + context: CallContext, + ) -> ir.Block: + """Generate the LLVM call operation and check if the call is successful.""" + with ir.InsertionPoint(current_block): + result = llvm.call( + result=self.i32_type, + callee=self.target_func, + callee_operands=call_operands, + op_bundle_sizes=[], + op_bundle_operands=[], + ) + return self.check_cuda_error(result, current_block, context) + + def insert_set_cuda_device(self, current_block: ir.Block, context: CallContext): + """Call the _cudaSetDevice function if we can find device id from tensor parameters.""" + + def find_cuda_device_index_from_params(): + for param in context.params: + if ( + isinstance(param, spec.Tensor) + and param.dlpack_device_type != tvm_ffi.DLDeviceType.kDLCPU + ): + return context.matched_var_binding[param.device_id] + return None + + cuda_device_index = find_cuda_device_index_from_params() + + if cuda_device_index is None: + return current_block + + with ir.InsertionPoint(current_block): + result = llvm.call( + result=self.i32_type, + callee="_cudaSetDevice", + callee_operands=[cuda_device_index], + op_bundle_sizes=[], + op_bundle_operands=[], + ) + + return self.check_cuda_error(result, current_block, context) + + def __call__(self, current_block: ir.Block, context: CallContext) -> ir.Block: + current_block = self.declare_extern_funcs(current_block, context) + current_block = self.insert_lazy_init_cuda(current_block, context) + current_block = self.append_unload_to_global_dtors(current_block, context) + current_block = self.insert_set_cuda_device(current_block, context) + current_block = super().__call__(current_block, context) + return current_block + + +class TVMFFIJitCompiledFunction(tvm_ffi.Function, CudaDialectJitCompiledFunction): + """TVM FFI Function that contains metadata of the compiled function and interface to the FFI layer. + + This function should not be directly used after + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # initialize the tvm_ffi.Function from the current execution engine + self._init_ffi_function() + + # use direct call to the tvm_ffi.Function.__call__ + # to avoid most of python overhead + __call__ = tvm_ffi.Function.__call__ + + def _init_ffi_function(self): + """Initialize the tvm_ffi.Function from the current execution engine. + + This function must be called at once during compilation time. + The reason why it is not called during init is because the original + flow may already created an execution engine and the function is not + guaranteed to be initialized at that time. + """ + if self.__chandle__() != 0: + raise DSLRuntimeError("TVM FFI function is already initialized") + # get the MLIR function pointer from the execution engine + tvm_ffi_function_ptr = self.engine.raw_lookup("__tvm_ffi_" + self.function_name) + tvm_ffi_function = tvm_ffi.Function.__from_mlir_packed_safe_call__( + tvm_ffi_function_ptr + ) + # move the handle from the tvm_ffi.Function to the current instance + self.__move_handle_from__(tvm_ffi_function) + + def to(self, device=None): + """TVM FFI function itself is already support all devices.""" + return self + + def run_compiled_program(self, exe_args: list[ir.Value]): + """Run the compiled program. This override is needed for implicit compile and execution.""" + return self.__call__(*exe_args) + + def export_to_c(self, object_file_path: str, function_name: str = None): + """Export the TVM FFI function to an object file. + + :param object_file_path: The path to the object file. + :param function_name: The name of the function to export. + """ + if function_name is not None and function_name != self.function_name: + mod = self.ir_module + rename_tvm_ffi_function(mod, self.function_name, function_name) + else: + mod = self.ir_module + + _execution_engine_extra.dump_object_file_pic( + mod, object_file_path, "__tvm_ffi_" + function_name, 2 + ) diff --git a/python/CuTeDSL/cutlass/pipeline/__init__.py b/python/CuTeDSL/cutlass/pipeline/__init__.py index 2f88819c..232d7b51 100644 --- a/python/CuTeDSL/cutlass/pipeline/__init__.py +++ b/python/CuTeDSL/cutlass/pipeline/__init__.py @@ -20,7 +20,9 @@ from .helpers import ( PipelineUserType, PipelineState, make_pipeline_state, + pipeline_init_arrive, pipeline_init_wait, + agent_sync, arrive, arrive_unaligned, wait, @@ -68,7 +70,9 @@ __all__ = [ "PipelineProducer", "PipelineConsumer", "make_pipeline_state", + "pipeline_init_arrive", "pipeline_init_wait", + "agent_sync", "arrive", "arrive_unaligned", "wait", diff --git a/python/CuTeDSL/cutlass/pipeline/helpers.py b/python/CuTeDSL/cutlass/pipeline/helpers.py index 480a0527..6787537c 100644 --- a/python/CuTeDSL/cutlass/pipeline/helpers.py +++ b/python/CuTeDSL/cutlass/pipeline/helpers.py @@ -16,7 +16,14 @@ from typing import Optional, Union import warnings import cutlass.cute as cute -from cutlass.cutlass_dsl import Boolean, Int32, Int64, if_generate +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 @@ -39,6 +46,11 @@ class Agent(enum.Enum): ThreadBlockCluster = enum.auto() +############################################################################## +# CooperativeGroup class +############################################################################## + + class CooperativeGroup: """ CooperativeGroup contains size and alignment restrictions for an Agent. @@ -74,6 +86,11 @@ class CooperativeGroup: self.agent = agent +############################################################################## +# PipelineOp class +############################################################################## + + class PipelineOp(enum.Enum): """ PipelineOp assigns an operation to an agent corresponding to a specific hardware feature. @@ -134,6 +151,11 @@ class SyncObject(ABC): pass +############################################################################## +# MbarrierArray class +############################################################################## + + class MbarrierArray(SyncObject): """ MbarrierArray implements an abstraction for an array of smem barriers. @@ -185,25 +207,33 @@ class MbarrierArray(SyncObject): return new_mbarrier_array # Mbarrier initialization - def mbarrier_init(self) -> None: + def mbarrier_init(self, *, loc=None, ip=None) -> None: """ Initializes an array of mbarriers using warp 0. """ def then_body(): for index in range(self.num_stages): - cute.arch.mbarrier_init(self.get_barrier(index), self.arrive_count) + cute.arch.mbarrier_init( + self.get_barrier(index, loc=loc, ip=ip), + self.arrive_count, + loc=loc, + ip=ip, + ) - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) + warp_idx = cute.arch.warp_idx(loc=loc, ip=ip) + warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip) - if_generate(warp_idx == 0, then_body) + if_generate(warp_idx == 0, then_body, loc=loc, ip=ip) def arrive( self, index: int, dst: int, cta_group: Optional[cute.nvgpu.tcgen05.CtaGroup] = None, + *, + loc=None, + ip=None, ) -> None: """Select the arrive corresponding to this MbarrierArray's PipelineOp. @@ -220,49 +250,75 @@ class MbarrierArray(SyncObject): :type cta_group: ``cute.nvgpu.tcgen05.CtaGroup``, optional """ if self.op_type is PipelineOp.AsyncThread: - self.arrive_mbarrier(index, dst) + self.arrive_mbarrier(index, dst, loc=loc, ip=ip) elif self.op_type is PipelineOp.TCGen05Mma: assert cta_group is not None, ( "Error: CTA group must be provided for TCGen05Mma." ) - self.arrive_tcgen05mma(index, dst, cta_group) + self.arrive_tcgen05mma(index, dst, cta_group, loc=loc, ip=ip) elif self.op_type in [PipelineOp.TmaLoad]: self.arrive_and_expect_tx(index, self.tx_count) elif self.op_type is PipelineOp.AsyncLoad: - self.arrive_cp_async_mbarrier(index) + self.arrive_cp_async_mbarrier(index, loc=loc, ip=ip) else: assert False, ( f"Error: MbarrierArray is not supported for PipelineOp: {_get_pipeline_op(self.op_type)}." ) - def arrive_mbarrier(self, index: int, dst_rank: Optional[int] = None) -> None: + def arrive_mbarrier( + self, index: int, dst_rank: Optional[int] = None, *, loc=None, ip=None + ) -> None: if dst_rank is None: - cute.arch.mbarrier_arrive(self.get_barrier(index)) + cute.arch.mbarrier_arrive( + self.get_barrier(index, loc=loc, ip=ip), loc=loc, ip=ip + ) else: - cute.arch.mbarrier_arrive(self.get_barrier(index), dst_rank) + cute.arch.mbarrier_arrive( + self.get_barrier(index, loc=loc, ip=ip), dst_rank, loc=loc, ip=ip + ) - def arrive_cp_async_mbarrier(self, index: int): - cute.arch.cp_async_mbarrier_arrive_noinc(self.get_barrier(index)) + def arrive_cp_async_mbarrier(self, index: int, *, loc=None, ip=None): + cute.arch.cp_async_mbarrier_arrive_noinc( + self.get_barrier(index, loc=loc, ip=ip), loc=loc, ip=ip + ) def arrive_tcgen05mma( - self, index: int, mask: Optional[int], cta_group: cute.nvgpu.tcgen05.CtaGroup + self, + index: int, + mask: Optional[int], + cta_group: cute.nvgpu.tcgen05.CtaGroup, + *, + loc=None, + ip=None, ) -> None: if mask is None: - with cute.arch.elect_one(): - cute.nvgpu.tcgen05.commit(self.get_barrier(index)) + with cute.arch.elect_one(loc=loc, ip=ip): + cute.nvgpu.tcgen05.commit( + self.get_barrier(index, loc=loc, ip=ip), loc=loc, ip=ip + ) else: - with cute.arch.elect_one(): - cute.nvgpu.tcgen05.commit(self.get_barrier(index), mask, cta_group) + with cute.arch.elect_one(loc=loc, ip=ip): + cute.nvgpu.tcgen05.commit( + self.get_barrier(index, loc=loc, ip=ip), + mask, + cta_group, + loc=loc, + ip=ip, + ) def arrive_and_expect_tx(self, index: int, tx_count: int) -> None: with cute.arch.elect_one(): cute.arch.mbarrier_arrive_and_expect_tx(self.get_barrier(index), tx_count) - def try_wait(self, index: int, phase: int) -> Boolean: - return cute.arch.mbarrier_try_wait(self.get_barrier(index), phase) + def try_wait(self, index: int, phase: int, *, loc=None, ip=None) -> Boolean: + return cute.arch.mbarrier_try_wait( + self.get_barrier(index, loc=loc, ip=ip), phase, loc=loc, ip=ip + ) - def wait(self, index: int, phase: int) -> None: - cute.arch.mbarrier_wait(self.get_barrier(index), phase) + def wait(self, index: int, phase: int, *, loc=None, ip=None) -> None: + cute.arch.mbarrier_wait( + self.get_barrier(index, loc=loc, ip=ip), phase, loc=loc, ip=ip + ) def arrive_and_wait( self, @@ -270,14 +326,17 @@ class MbarrierArray(SyncObject): phase: int, dst: int, cta_group: Optional[cute.nvgpu.tcgen05.CtaGroup] = None, + *, + loc=None, + ip=None, ) -> None: - arrive(index, dst, cta_group) - wait(index, phase) + arrive(index, dst, cta_group, loc=loc, ip=ip) + wait(index, phase, loc=loc, ip=ip) - def arrive_and_drop(self) -> None: + def arrive_and_drop(self, *, loc=None, ip=None) -> None: raise NotImplementedError("Error: Not yet supported.") - def get_barrier(self, index: int) -> cute.Pointer: + def get_barrier(self, index: int, *, loc=None, ip=None) -> cute.Pointer: return self.mbarrier_base + index def max(self) -> int: @@ -294,6 +353,11 @@ class MbarrierArray(SyncObject): ) +############################################################################## +# NamedBarrier class +############################################################################## + + @dataclass(frozen=True) class NamedBarrier(SyncObject): """ @@ -314,16 +378,21 @@ class NamedBarrier(SyncObject): "NamedBarrier ID 0 is by other driver APIs (i.e. sync_threads()) and should not be used." ) - def arrive(self) -> None: + @dsl_user_op + def arrive(self, *, loc=None, ip=None) -> None: """ The aligned flavor of arrive is used when all threads in the CTA will execute the same instruction. See PTX documentation. """ cute.arch.barrier_arrive( - barrier_id=self.barrier_id, number_of_threads=self.num_threads + barrier_id=self.barrier_id, + number_of_threads=self.num_threads, + loc=loc, + ip=ip, ) - def arrive_unaligned(self) -> None: + @dsl_user_op + def arrive_unaligned(self, *, loc=None, ip=None) -> None: """ The unaligned flavor of arrive can be used with an arbitrary number of threads in the CTA. """ @@ -337,7 +406,8 @@ class NamedBarrier(SyncObject): asm_dialect=llvm.AsmDialect.AD_ATT, ) - def wait(self) -> None: + @dsl_user_op + def wait(self, *, loc=None, ip=None) -> 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 @@ -348,7 +418,7 @@ class NamedBarrier(SyncObject): warnings.warn( "NamedBarrier wait also arrives on the barrier. Routing call to NamedBarrier.arrive_and_wait()." ) - self.arrive_and_wait() + self.arrive_and_wait(loc=loc, ip=ip) def wait_unaligned(self) -> None: llvm.inline_asm( @@ -361,18 +431,22 @@ class NamedBarrier(SyncObject): asm_dialect=llvm.AsmDialect.AD_ATT, ) - def arrive_and_wait(self) -> None: + @dsl_user_op + def arrive_and_wait(self, *, loc=None, ip=None) -> None: cute.arch.barrier( - barrier_id=self.barrier_id, number_of_threads=self.num_threads + barrier_id=self.barrier_id, + number_of_threads=self.num_threads, + loc=loc, + ip=ip, ) - def arrive_and_drop(self) -> None: + def arrive_and_drop(self, *, loc=None, ip=None) -> None: raise NotImplementedError("Error: Not supported.") - def sync(self) -> None: + def sync(self, *, loc=None, ip=None) -> None: self.arrive_and_wait() - def get_barrier(self) -> int: + def get_barrier(self, *, loc=None, ip=None) -> int: return self.barrier_id def max(self) -> int: @@ -380,6 +454,11 @@ class NamedBarrier(SyncObject): return 4095 +############################################################################## +# TmaStoreFence class +############################################################################## + + class TmaStoreFence(SyncObject): """ TmaStoreFence is used for a multi-stage epilogue buffer. @@ -391,21 +470,27 @@ class TmaStoreFence(SyncObject): self.num_stages = num_stages - def arrive(self) -> None: - cute.arch.cp_async_bulk_commit_group() + @dsl_user_op + def arrive(self, *, loc=None, ip=None) -> None: + cute.arch.cp_async_bulk_commit_group(loc=loc, ip=ip) - def wait(self) -> None: - cute.arch.cp_async_bulk_wait_group(self.num_stages - 1, read=True) + @dsl_user_op + def wait(self, *, loc=None, ip=None) -> None: + cute.arch.cp_async_bulk_wait_group( + self.num_stages - 1, read=True, loc=loc, ip=ip + ) - def arrive_and_wait(self) -> None: - self.arrive() - self.wait() + @dsl_user_op + def arrive_and_wait(self, *, loc=None, ip=None) -> None: + self.arrive(loc=loc, ip=ip) + self.wait(loc=loc, ip=ip) - def arrive_and_drop(self) -> None: + @dsl_user_op + def arrive_and_drop(self, *, loc=None, ip=None) -> None: raise NotImplementedError("Error: Not supported.") # TmaStoreFence doesn't have mbarriers - def get_barrier(self) -> None: + def get_barrier(self, *, loc=None, ip=None) -> None: assert False, ( "Error: TmaStoreFence doesn't use mbarriers and cannot return a barrier." ) @@ -413,12 +498,13 @@ class TmaStoreFence(SyncObject): def max(self) -> None: raise NotImplementedError("Error: Not supported.") - def tail(self) -> None: - cute.arch.cp_async_bulk_wait_group(0, read=True) + @dsl_user_op + def tail(self, *, loc=None, ip=None) -> None: + cute.arch.cp_async_bulk_wait_group(0, read=True, loc=loc, ip=ip) ############################################################################## -# PipelineState class +# PipelineUserType class ############################################################################## @@ -427,6 +513,11 @@ class PipelineUserType(enum.Enum): Consumer = enum.auto() +############################################################################## +# PipelineState class +############################################################################## + + class PipelineState: """ Pipeline state contains an index and phase bit corresponding to the current position in the circular buffer. @@ -460,12 +551,13 @@ class PipelineState: def reset_count(self): self._count = Int32(0) - def advance(self): + @dsl_user_op + def advance(self, *, loc=None, ip=None): self._index += 1 self._count += 1 def then_body(index, phase): - new_index = Int32(0) + new_index = Int32(0, loc=loc, ip=ip) new_phase = phase ^ 1 return new_index, new_phase @@ -478,14 +570,17 @@ class PipelineState: else_body, [self.index, self.phase], [Int32, Int32], + loc=loc, + ip=ip, ) - def reverse(self): + @dsl_user_op + def reverse(self, *, loc=None, ip=None): self._index -= 1 self._count -= 1 def then_body(index, phase): - new_index = Int32(self.stages - 1) + new_index = Int32(self.stages - 1, loc=loc, ip=ip) new_phase = phase ^ 1 return new_index, new_phase @@ -498,6 +593,8 @@ class PipelineState: else_body, [self.index, self.phase], [Int32, Int32], + loc=loc, + ip=ip, ) def __get_mlir_types__(self): @@ -516,23 +613,23 @@ class PipelineState: ) -def make_pipeline_state(type: PipelineUserType, stages: int): +def make_pipeline_state(type: PipelineUserType, stages: int, *, loc=None, ip=None): """ Creates a pipeline state. Producers are assumed to start with an empty buffer and have a flipped phase bit of 1. """ if type is PipelineUserType.Producer: return PipelineState( stages, - Int32(0), - Int32(0), - Int32(1), + Int32(0, loc=loc, ip=ip), + Int32(0, loc=loc, ip=ip), + Int32(1, loc=loc, ip=ip), ) elif type is PipelineUserType.Consumer: return PipelineState( stages, - Int32(0), - Int32(0), - Int32(0), + Int32(0, loc=loc, ip=ip), + Int32(0, loc=loc, ip=ip), + Int32(0, loc=loc, ip=ip), ) else: assert False, ( @@ -545,31 +642,65 @@ def make_pipeline_state(type: PipelineUserType, stages: int): ############################################################################## -def pipeline_init_wait(cta_layout_vmnk: Optional[cute.Layout] = None): +def pipeline_init_arrive( + cluster_shape_mn: Optional[cute.Layout] = None, + is_relaxed: bool = False, + *, + loc=None, + ip=None, +): """ - Fences the mbarrier init and syncs the threadblock or cluster + Fences the mbarrier_init and sends an arrive if using clusters. """ - cute.arch.mbarrier_init_fence() - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1: + # If using clusters, send nonblocking arrives. Otherwise, do nothing + # because sync_threads() doesn't have a nonblocking arrive. + + cute.arch.mbarrier_init_fence(loc=loc, ip=ip) + + if cluster_shape_mn is not None and cute.size(cluster_shape_mn, loc=loc, ip=ip) > 1: + if is_relaxed: + # Fences memory operations issued before the arrive + cute.arch.cluster_arrive_relaxed(loc=loc, ip=ip) + else: + # Skips the memory barrier + cute.arch.cluster_arrive(loc=loc, ip=ip) + + +def pipeline_init_wait( + cluster_shape_mn: Optional[cute.Layout] = None, *, loc=None, ip=None +): + """ + Syncs the threadblock or cluster + """ + + if cluster_shape_mn is None or cute.size(cluster_shape_mn, loc=loc, ip=ip) == 1: # If not using clusters, sync the threadblock - _sync(Agent.ThreadBlock) + agent_sync(Agent.ThreadBlock, loc=loc, ip=ip) else: - # If using clusters, sync the cluster - _sync(Agent.ThreadBlockCluster) + # If using clusters, wait on the cluster + cute.arch.cluster_wait(loc=loc, ip=ip) -def _sync(group: Agent): +def _sync(group: Agent, is_relaxed: bool = False, *, loc=None, ip=None): + warnings.warn("_sync is deprecated. Please use agent_sync instead.") + agent_sync(group, is_relaxed, loc=loc, ip=ip) + + +def agent_sync(group: Agent, is_relaxed: bool = False, *, loc=None, ip=None): """ Syncs all threads within an agent. """ if group is Agent.Thread: raise NotImplementedError("Error: Not supported.") elif group is Agent.ThreadBlock: - cute.arch.sync_threads() + cute.arch.sync_threads(loc=loc, ip=ip) elif group is Agent.ThreadBlockCluster: - cute.arch.cluster_arrive() - cute.arch.cluster_wait() + if is_relaxed: + cute.arch.cluster_arrive_relaxed(loc=loc, ip=ip) + else: + cute.arch.cluster_arrive(loc=loc, ip=ip) + cute.arch.cluster_wait(loc=loc, ip=ip) else: assert False, ( "Error: No explicit sync instruction exists. Please use barriers (named / mbarrier) instead." @@ -589,15 +720,17 @@ def _mbarrier_i64_to_ptr(val: Int64) -> cute.Pointer: # NamedBarrier free functions -def arrive(barrier_id: int, num_threads: int): +def arrive(barrier_id: int, num_threads: int, *, loc=None, ip=None): """ The aligned flavor of arrive is used when all threads in the CTA will execute the same instruction. See PTX documentation. """ - cute.arch.barrier_arrive(barrier_id=barrier_id, number_of_threads=num_threads) + cute.arch.barrier_arrive( + barrier_id=barrier_id, number_of_threads=num_threads, loc=loc, ip=ip + ) -def arrive_unaligned(barrier_id: int, num_threads: int): +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. """ @@ -623,10 +756,10 @@ def wait(barrier_id: int, num_threads: int): warnings.warn( "NamedBarrier wait also arrives on the barrier. Routing call to NamedBarrier.arrive_and_wait()." ) - arrive_and_wait() + arrive_and_wait(loc=loc, ip=ip) -def wait_unaligned(barrier_id: int, num_threads: int): +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()." ) @@ -641,9 +774,11 @@ def wait_unaligned(barrier_id: int, num_threads: int): ) -def arrive_and_wait(barrier_id: int, num_threads: int): - cute.arch.barrier(barrier_id=barrier_id, number_of_threads=num_threads) +def arrive_and_wait(barrier_id: int, num_threads: int, *, loc=None, ip=None): + cute.arch.barrier( + barrier_id=barrier_id, number_of_threads=num_threads, loc=loc, ip=ip + ) -def sync(barrier_id: int = 0): - cute.arch.barrier(barrier_id=barrier_id) +def sync(barrier_id: int = 0, *, loc=None, ip=None): + cute.arch.barrier(barrier_id=barrier_id, loc=loc, ip=ip) diff --git a/python/CuTeDSL/cutlass/pipeline/sm100.py b/python/CuTeDSL/cutlass/pipeline/sm100.py index 60cb83e2..180da44d 100644 --- a/python/CuTeDSL/cutlass/pipeline/sm100.py +++ b/python/CuTeDSL/cutlass/pipeline/sm100.py @@ -17,11 +17,12 @@ import cutlass.cute as cute from cutlass.cutlass_dsl import Boolean, if_generate from cutlass.pipeline import ( + Agent, CooperativeGroup, PipelineOp, PipelineState, - pipeline_init_wait, PipelineAsync, + agent_sync, ) ############################################################################## @@ -106,23 +107,27 @@ class PipelineTmaUmma(PipelineAsync): barrier_storage: cute.Pointer = None, cta_layout_vmnk: Optional[cute.Layout] = None, mcast_mode_mn: tuple[int, int] = (1, 1), + defer_sync: bool = False, ): - """ - This helper function computes any necessary attributes and returns an instance of PipelineTmaUmma. - :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers - :type barrier_storage: cute.Pointer + """Creates and initializes a new PipelineTmaUmma instance. + :param num_stages: Number of buffer stages for this pipeline - :type num_stages: Int32 - :param producer_group: `CooperativeGroup` for the producer agent + :type num_stages: int + :param producer_group: CooperativeGroup for the producer agent :type producer_group: CooperativeGroup - :param consumer_group: `CooperativeGroup` for the consumer agent + :param consumer_group: CooperativeGroup for the consumer agent :type consumer_group: CooperativeGroup :param tx_count: Number of bytes expected to be written to the transaction barrier for one stage :type tx_count: int + :param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers + :type barrier_storage: cute.Pointer, optional :param cta_layout_vmnk: Layout of the cluster shape - :type cta_layout_vmnk: cute.Layout | None - :param mcast_mode_mn: Tuple of two integers, specifying whether mcast is enabled for the m and n modes. At least one of the two integers must be 1. - :type mcast_mode_mn: tuple[int, int] + :type cta_layout_vmnk: cute.Layout, optional + :param mcast_mode_mn: Tuple specifying multicast modes for m and n dimensions (each 0 or 1) + :type mcast_mode_mn: tuple[int, int], optional + :raises ValueError: If barrier_storage is not a cute.Pointer instance + :return: A new PipelineTmaUmma instance configured with the provided parameters + :rtype: PipelineTmaUmma """ if not isinstance(barrier_storage, cute.Pointer): raise ValueError( @@ -161,7 +166,12 @@ class PipelineTmaUmma(PipelineAsync): consumer_mask = producer_mask - pipeline_init_wait(cta_layout_vmnk) + 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: + agent_sync(Agent.ThreadBlockCluster, is_relaxed=True) return PipelineTmaUmma( sync_object_full, @@ -262,19 +272,23 @@ class PipelineAsyncUmma(PipelineAsync): consumer_group: CooperativeGroup, barrier_storage: cute.Pointer = None, cta_layout_vmnk: Optional[cute.Layout] = None, + defer_sync: bool = False, ): - """ - This helper function computes any necessary attributes and returns an instance of PipelineAsyncUmma. - :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers - :type barrier_storage: cute.Pointer + """Creates and initializes a new PipelineAsyncUmma instance. + :param num_stages: Number of buffer stages for this pipeline - :type num_stages: Int32 - :param producer_group: `CooperativeGroup` for the producer agent + :type num_stages: int + :param producer_group: CooperativeGroup for the producer agent :type producer_group: CooperativeGroup - :param consumer_group: `CooperativeGroup` for the consumer agent + :param consumer_group: CooperativeGroup for the consumer agent :type consumer_group: CooperativeGroup + :param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers + :type barrier_storage: cute.Pointer, optional :param cta_layout_vmnk: Layout of the cluster shape - :type cta_layout_vmnk: cute.Layout | None + :type cta_layout_vmnk: cute.Layout, optional + :raises ValueError: If barrier_storage is not a cute.Pointer instance + :return: A new PipelineAsyncUmma instance configured with the provided parameters + :rtype: PipelineAsyncUmma """ if not isinstance(barrier_storage, cute.Pointer): raise ValueError( @@ -315,7 +329,12 @@ class PipelineAsyncUmma(PipelineAsync): # consumer needs to get the mask to signal consumer_mask = PipelineAsyncUmma._compute_peer_cta_mask(cta_layout_vmnk) - pipeline_init_wait(cta_layout_vmnk) + 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: + agent_sync(Agent.ThreadBlockCluster, is_relaxed=True) return PipelineAsyncUmma( sync_object_full, @@ -372,19 +391,23 @@ class PipelineUmmaAsync(PipelineAsync): consumer_group: CooperativeGroup, barrier_storage: cute.Pointer = None, cta_layout_vmnk: Optional[cute.Layout] = None, + defer_sync: bool = False, ): - """ - This helper function computes any necessary attributes and returns an instance of PipelineUmmaAsync. - :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers - :type barrier_storage: cute.Pointer + """Creates an instance of PipelineUmmaAsync with computed attributes. + :param num_stages: Number of buffer stages for this pipeline - :type num_stages: Int32 - :param producer_group: `CooperativeGroup` for the producer agent + :type num_stages: int + :param producer_group: ``CooperativeGroup`` for the producer agent :type producer_group: CooperativeGroup - :param consumer_group: `CooperativeGroup` for the consumer agent + :param consumer_group: ``CooperativeGroup`` for the consumer agent :type consumer_group: CooperativeGroup + :param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers + :type barrier_storage: cute.Pointer, optional :param cta_layout_vmnk: Layout of the cluster shape - :type cta_layout_vmnk: cute.Layout | None + :type cta_layout_vmnk: cute.Layout, optional + :raises ValueError: If barrier_storage is not a cute.Pointer instance + :return: New instance of ``PipelineUmmaAsync`` + :rtype: PipelineUmmaAsync """ if not isinstance(barrier_storage, cute.Pointer): raise ValueError( @@ -411,7 +434,7 @@ class PipelineUmmaAsync(PipelineAsync): producer_mask = PipelineUmmaAsync._compute_tmem_sync_mask(cta_layout_vmnk) if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1: - # Set mask to None if not using 2CTA intructions + # Set mask to None if not using 2CTA instructions consumer_mask = None else: consumer_mask = PipelineUmmaAsync._compute_peer_cta_rank() @@ -422,7 +445,12 @@ class PipelineUmmaAsync(PipelineAsync): else cute.nvgpu.tcgen05.CtaGroup.TWO ) - pipeline_init_wait(cta_layout_vmnk) + 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: + agent_sync(Agent.ThreadBlockCluster, is_relaxed=True) return PipelineUmmaAsync( sync_object_full, diff --git a/python/CuTeDSL/cutlass/pipeline/sm90.py b/python/CuTeDSL/cutlass/pipeline/sm90.py index 8e926354..b7cef128 100644 --- a/python/CuTeDSL/cutlass/pipeline/sm90.py +++ b/python/CuTeDSL/cutlass/pipeline/sm90.py @@ -13,8 +13,10 @@ from dataclasses import dataclass from typing import Optional import cutlass.cute as cute -from cutlass.cutlass_dsl import Boolean, Int32, if_generate +from cutlass.cutlass_dsl import Boolean, Int32, if_generate, dsl_user_op + from cutlass.pipeline import ( + Agent, CooperativeGroup, MbarrierArray, PipelineOp, @@ -23,7 +25,7 @@ from cutlass.pipeline import ( SyncObject, TmaStoreFence, make_pipeline_state, - pipeline_init_wait, + agent_sync, ) ############################################################################## @@ -157,6 +159,7 @@ class PipelineAsync: barrier_storage: cute.Pointer = None, producer_mask: Int32 = None, consumer_mask: Int32 = None, + defer_sync: bool = False, ): """Creates and initializes a new PipelineAsync instance. @@ -167,17 +170,17 @@ class PipelineAsync: :type barrier_storage: cute.Pointer :param num_stages: Number of buffer stages for this pipeline :type num_stages: int - :param producer_group: `CooperativeGroup` for the producer agent + :param producer_group: ``CooperativeGroup`` for the producer agent :type producer_group: CooperativeGroup - :param consumer_group: `CooperativeGroup` for the consumer agent + :param consumer_group: ``CooperativeGroup`` for the consumer agent :type consumer_group: CooperativeGroup - :param producer_mask: Mask for signaling arrives for the producer agent, defaults to ``None`` + :param producer_mask: Mask for signaling arrives for the producer agent :type producer_mask: Int32, optional - :param consumer_mask: Mask for signaling arrives for the consumer agent, defaults to ``None`` + :param consumer_mask: Mask for signaling arrives for the consumer agent :type consumer_mask: Int32, optional - :return: A new PipelineAsync instance - :rtype: PipelineAsync :raises ValueError: If barrier_storage is not a cute.Pointer instance + :return: A new ``PipelineAsync`` instance + :rtype: PipelineAsync """ if not isinstance(barrier_storage, cute.Pointer): raise ValueError( @@ -197,7 +200,9 @@ class PipelineAsync: barrier_storage.align(min_align=8) + num_stages, num_stages, consumer ) - pipeline_init_wait() + if not defer_sync: + cute.arch.mbarrier_init_fence() + agent_sync(Agent.ThreadBlock) return PipelineAsync( sync_object_full, @@ -280,21 +285,24 @@ class PipelineCpAsync(PipelineAsync): consumer_group: CooperativeGroup, producer_mask: Int32 = None, consumer_mask: Int32 = None, + defer_sync: bool = False, ): - """ - This helper function computes any necessary attributes and returns an instance of PipelineAsync. - :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers + """Helper function that computes necessary attributes and returns a ``PipelineCpAsync`` instance. + + :param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers :type barrier_storage: cute.Pointer :param num_stages: Number of buffer stages for this pipeline :type num_stages: Int32 - :param producer_group: CooperativeGroup for the producer agent + :param producer_group: ``CooperativeGroup`` for the producer agent :type producer_group: CooperativeGroup - :param consumer_group: CooperativeGroup for the consumer agent + :param consumer_group: ``CooperativeGroup`` for the consumer agent :type consumer_group: CooperativeGroup - :param producer_mask: Mask for signaling arrives for the producer agent - :type producer_mask: Int32 | None - :param consumer_mask: Mask for signaling arrives for the consumer agent - :type consumer_mask: Int32 | None + :param producer_mask: Mask for signaling arrives for the producer agent, defaults to None + :type producer_mask: Int32, optional + :param consumer_mask: Mask for signaling arrives for the consumer agent, defaults to None + :type consumer_mask: Int32, optional + :return: A new ``PipelineCpAsync`` instance configured with the provided parameters + :rtype: PipelineCpAsync """ producer_type = PipelineOp.AsyncLoad consumer_type = PipelineOp.AsyncThread @@ -309,7 +317,9 @@ class PipelineCpAsync(PipelineAsync): barrier_storage.align(min_align=8) + num_stages, num_stages, consumer ) - pipeline_init_wait() + if not defer_sync: + cute.arch.mbarrier_init_fence() + agent_sync(Agent.ThreadBlock) return PipelineCpAsync( sync_object_array_full, @@ -331,11 +341,22 @@ class PipelineTmaAsync(PipelineAsync): @staticmethod @cute.jit def init_empty_barrier_arrive_signal( - cta_layout_vmnk: cute.Layout, tidx: Int32, mcast_mode_mn: tuple[int, int] + cta_layout_vmnk: cute.Layout, tidx: Int32, mcast_mode_mn: tuple[int, int] = (1, 1) ): - """ - Initialize the empty barrier arrive signal - This function returns the destination cta rank and a boolean indicating if the signalling thread is the same as the current thread + """Initialize the empty barrier arrive signal. + + This function determines which threads should signal empty barrier arrives based on the cluster layout + and multicast modes. It returns the destination CTA rank and whether the current thread should signal. + + :param cta_layout_vmnk: Layout describing the cluster shape and CTA arrangement + :type cta_layout_vmnk: cute.Layout + :param tidx: Thread index within the warp + :type tidx: Int32 + :param mcast_mode_mn: Tuple specifying multicast modes for m and n dimensions (each 0 or 1), defaults to (1,1) + :type mcast_mode_mn: tuple[int, int] + :raises ``AssertionError``: If both multicast modes are disabled (0,0) + :return: Tuple containing destination CTA rank and boolean indicating if current thread signals + :rtype: tuple[Int32, Boolean] """ # Logic to optimally schedule Empty Arrives cluster_shape_vmnk = cta_layout_vmnk.shape @@ -384,25 +405,29 @@ class PipelineTmaAsync(PipelineAsync): cta_layout_vmnk: Optional[cute.Layout] = None, tidx: Optional[Int32] = None, mcast_mode_mn: tuple[int, int] = (1, 1), + defer_sync: bool = False, ): - """ - This helper function computes any necessary attributes and returns an instance of PipelineTmaAsync. - :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers - :type barrier_storage: cute.Pointer + """Create a new ``PipelineTmaAsync`` instance. + :param num_stages: Number of buffer stages for this pipeline - :type num_stages: Int32 - :param producer_group: `CooperativeGroup` for the producer agent + :type num_stages: int + :param producer_group: ``CooperativeGroup`` for the producer agent :type producer_group: CooperativeGroup - :param consumer_group: `CooperativeGroup` for the consumer agent + :param consumer_group: ``CooperativeGroup`` for the consumer agent :type consumer_group: CooperativeGroup :param tx_count: Number of bytes expected to be written to the transaction barrier for one stage :type tx_count: int - :param cta_layout_vmnk: Layout of the cluster shape - :type cta_layout_vmnk: cute.Layout | None - :param tidx: thread index to consumer async threads - :type tidx: Int32 | None - :param mcast_mode_mn: Tuple of two integers, specifying whether mcast is enabled for the m and n modes. At least one of the two integers must be 1. - :type mcast_mode_mn: tuple[int, int] + :param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers, defaults to None + :type barrier_storage: cute.Pointer, optional + :param cta_layout_vmnk: Layout of the cluster shape, defaults to None + :type cta_layout_vmnk: cute.Layout, optional + :param tidx: Thread index to consumer async threads, defaults to None + :type tidx: Int32, optional + :param mcast_mode_mn: Tuple specifying multicast modes for m and n dimensions (each 0 or 1), defaults to (1,1) + :type mcast_mode_mn: tuple[int, int], optional + :raises ValueError: If barrier_storage is not a cute.Pointer instance + :return: New ``PipelineTmaAsync`` instance + :rtype: PipelineTmaAsync """ if not isinstance(barrier_storage, cute.Pointer): raise ValueError( @@ -438,7 +463,12 @@ class PipelineTmaAsync(PipelineAsync): producer_mask = None - pipeline_init_wait(cta_layout_vmnk) + 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: + agent_sync(Agent.ThreadBlockCluster, is_relaxed=True) return PipelineTmaAsync( sync_object_full, @@ -449,17 +479,25 @@ class PipelineTmaAsync(PipelineAsync): is_signalling_thread, ) + @dsl_user_op def producer_acquire( - self, state: PipelineState, try_acquire_token: Optional[Boolean] = None + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + *, + loc=None, + ip=None, ): """ TMA producer commit conditionally waits on buffer empty and sets the transaction barrier. """ if_generate( try_acquire_token is None or try_acquire_token == 0, - lambda: self.sync_object_empty.wait(state.index, state.phase), + lambda: self.sync_object_empty.wait( + state.index, state.phase, loc=loc, ip=ip + ), ) - self.sync_object_full.arrive(state.index, self.producer_mask) + self.sync_object_full.arrive(state.index, self.producer_mask, loc=loc, ip=ip) def producer_commit(self, state: PipelineState): """ @@ -467,13 +505,16 @@ class PipelineTmaAsync(PipelineAsync): """ pass - def consumer_release(self, state: PipelineState): + @dsl_user_op + def consumer_release(self, state: PipelineState, *, loc=None, ip=None): """ TMA consumer release conditionally signals the empty buffer to the producer. """ if_generate( self.is_signalling_thread, - lambda: self.sync_object_empty.arrive(state.index, self.consumer_mask), + lambda: self.sync_object_empty.arrive( + state.index, self.consumer_mask, loc=loc, ip=ip + ), ) @@ -498,23 +539,29 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): tx_count: int, barrier_storage: cute.Pointer = None, cta_layout_vmnk: Optional[cute.Layout] = None, + defer_sync: bool = False, ): - """ - This helper function computes any necessary attributes and returns an instance of PipelineTmaMultiConsumersAsync. - :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers + """Creates an instance of PipelineTmaMultiConsumersAsync with computed attributes. + + :param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers :type barrier_storage: cute.Pointer :param num_stages: Number of buffer stages for this pipeline - :type num_stages: Int32 - :param producer_group: `CooperativeGroup` for the producer agent + :type num_stages: int + :param producer_group: ``CooperativeGroup`` for the producer agent :type producer_group: CooperativeGroup - :param consumer_group_umma: `CooperativeGroup` for the UMMA consumer agent + :param consumer_group_umma: ``CooperativeGroup`` for the UMMA consumer agent :type consumer_group_umma: CooperativeGroup - :param consumer_group_async: `CooperativeGroup` for the AsyncThread consumer agent + :param consumer_group_async: ``CooperativeGroup`` for the AsyncThread consumer agent :type consumer_group_async: CooperativeGroup :param tx_count: Number of bytes expected to be written to the transaction barrier for one stage :type tx_count: int - :param cta_layout_vmnk: Layout of the cluster shape - :type cta_layout_vmnk: cute.Layout | None + :param cta_layout_vmnk: Layout of the cluster shape, defaults to None + :type cta_layout_vmnk: Optional[cute.Layout] + :raises ValueError: If ``barrier_storage`` is not a ``cute.Pointer`` instance + :raises ValueError: If ``UMMA`` and ``AsyncThread`` consumer groups are not the same agent + :raises ValueError: If ``cta_layout_vmnk`` size is not 1 + :return: New instance of ``PipelineTmaMultiConsumersAsync`` + :rtype: PipelineTmaMultiConsumersAsync """ if not isinstance(barrier_storage, cute.Pointer): raise ValueError( @@ -560,7 +607,7 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): # No mcast mask if not using clusters producer_mask = None consumer_mask = None - # All threadblocks are leaders if not using clusters + # All thread-blocks are leaders if not using clusters is_leader_cta = True cta_group = ( cute.nvgpu.tcgen05.CtaGroup.ONE @@ -568,7 +615,11 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): else cute.nvgpu.tcgen05.CtaGroup.TWO ) - pipeline_init_wait(cta_layout_vmnk) + if not defer_sync: + if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1: + agent_sync(Agent.ThreadBlock) + else: + agent_sync(Agent.ThreadBlockCluster, is_relaxed=True) return PipelineTmaMultiConsumersAsync( sync_object_full, @@ -626,14 +677,15 @@ class PipelineTmaStore(PipelineAsync): num_stages: int, producer_group: CooperativeGroup, ): - """ - This helper function computes any necessary attributes and returns an instance of PipelineTmaStore. - :param num_stages: Number of buffer stages for this pipeline - :type num_stages: Int32 - :param producer_group: `CooperativeGroup` for the producer agent - :type producer_group: CooperativeGroup - """ + """This helper function computes any necessary attributes and returns an instance of ``PipelineTmaStore``. + :param num_stages: Number of buffer stages for this pipeline + :type num_stages: int + :param producer_group: ``CooperativeGroup`` for the producer agent + :type producer_group: CooperativeGroup + :return: A new ``PipelineTmaStore`` instance + :rtype: PipelineTmaStore + """ producer_type = PipelineOp.TmaStore producer = (producer_type, producer_group) @@ -642,20 +694,25 @@ class PipelineTmaStore(PipelineAsync): return PipelineTmaStore(sync_object_full, None, num_stages, None, None) - def producer_acquire(self): - self.sync_object_full.wait() + @dsl_user_op + def producer_acquire(self, *, loc=None, ip=None): + self.sync_object_full.wait(loc=loc, ip=ip) - def producer_commit(self): - self.sync_object_full.arrive() + @dsl_user_op + def producer_commit(self, *, loc=None, ip=None): + self.sync_object_full.arrive(loc=loc, ip=ip) - def consumer_wait(self): + @dsl_user_op + def consumer_wait(self, *, loc=None, ip=None): assert False, "Error: PipelineTmaStore does not have a consumer agent." - def consumer_release(self): + @dsl_user_op + def consumer_release(self, *, loc=None, ip=None): assert False, "Error: PipelineTmaStore does not have a consumer agent." - def producer_tail(self): - self.sync_object_full.tail() + @dsl_user_op + def producer_tail(self, *, loc=None, ip=None): + self.sync_object_full.tail(loc=loc, ip=ip) @dataclass(frozen=True) @@ -706,6 +763,7 @@ class PipelineOrder: length: int, group_id: int, producer_group: CooperativeGroup, + defer_sync: bool = False, ): if not isinstance(barrier_storage, cute.Pointer): raise ValueError( @@ -722,7 +780,9 @@ class PipelineOrder: barrier_storage.align(min_align=8), num_stages, producer ) - pipeline_init_wait() + if not defer_sync: + cute.arch.mbarrier_init_fence() + agent_sync(Agent.ThreadBlock) return PipelineOrder( sync_object_full, @@ -740,16 +800,20 @@ class PipelineOrder: def get_barrier_for_current_stage_idx(self, group_id): return self.state.index * self.length + group_id - def arrive(self): + @dsl_user_op + def arrive(self, *, loc=None, ip=None): signalling_id = (self.group_id + 1) % self.length idx = self.get_barrier_for_current_stage_idx(signalling_id) - cute.arch.mbarrier_arrive(self.sync_object_full.get_barrier(idx)) - self.state.advance() + cute.arch.mbarrier_arrive( + self.sync_object_full.get_barrier(idx), loc=loc, ip=ip + ) + self.state.advance(loc=loc, ip=ip) - def wait(self): + @dsl_user_op + def wait(self, *, loc=None, ip=None): idx = self.get_barrier_for_current_stage_idx(self.group_id) cute.arch.mbarrier_wait( - self.sync_object_full.get_barrier(idx), self.state.phase + self.sync_object_full.get_barrier(idx), self.state.phase, loc=loc, ip=ip ) @@ -881,6 +945,10 @@ class PipelineProducer: self.__state = state self.__group = group + def clone(self): + """Create a new Producer instance with the same state.""" + return PipelineProducer(self.__pipeline, self.__state.clone(), self.__group) + def reset(self): """Reset the count of how many handles this producer has committed.""" self.__state.reset_count() @@ -1053,6 +1121,10 @@ class PipelineConsumer: self.__group = group self.__state = state + def clone(self): + """Create a new Consumer instance with the same state.""" + return PipelineConsumer(self.__pipeline, self.__state.clone(), self.__group) + def reset(self): """Reset the count of how many handles this consumer has consumed.""" self.__state.reset_count() diff --git a/python/CuTeDSL/cutlass/torch.py b/python/CuTeDSL/cutlass/torch.py index 09c30609..ac3fe2e2 100644 --- a/python/CuTeDSL/cutlass/torch.py +++ b/python/CuTeDSL/cutlass/torch.py @@ -47,6 +47,11 @@ def dtype(ty: Type[Numeric]): Float8E4M3FN: torch.float8_e4m3fn, Float8E4M3B11FNUZ: torch.float8_e4m3fnuz, } + + # float8_e8m0fnu is introduced in latest version of torch + if hasattr(torch, "float8_e8m0fnu"): + torch_type_map[Float8E8M0FNU] = torch.float8_e8m0fnu + if torch_dtype is None: torch_dtype = torch_type_map.get(ty) diff --git a/python/CuTeDSL/cutlass/utils/__init__.py b/python/CuTeDSL/cutlass/utils/__init__.py index 024a4096..3cee6d55 100644 --- a/python/CuTeDSL/cutlass/utils/__init__.py +++ b/python/CuTeDSL/cutlass/utils/__init__.py @@ -13,6 +13,7 @@ from .static_persistent_tile_scheduler import ( WorkTileInfo, PersistentTileSchedulerParams, StaticPersistentTileScheduler, + StaticPersistentRuntimeTileScheduler, ) from .hardware_info import ( @@ -65,21 +66,27 @@ from .tmem_allocator import TmemAllocator from .layout import LayoutEnum -from .distributed_helpers import ( - spin_lock_wait, - spin_lock_multimem_arrive, - multimem_ld_reduce_8xf16, - multimem_ld_reduce_4xf32, - multimem_ld_reduce_8xbf16, - multimem_ld_reduce_16xe4m3, - multimem_ld_reduce_16xe5m2, - multimem_st_4xb32, - sm_wise_inter_gpu_multimem_barrier, +from . import distributed + +from .mixed_input_helpers import ( + TransformMode, + scale_tma_partition, + transform_partition, + scale_partition, + get_gmem_layout_scale, + get_smem_layout_scale, + compute_smem_layout, + get_transform_a_source, + get_tma_atom_kind, + get_copy_atom_a_transform, + is_valid_scale_granularity, + get_divisibility, ) from . import hopper_helpers as sm90 from . import blackwell_helpers as sm100 +from .print_latex import print_latex, print_latex_tv __all__ = [ "get_smem_capacity_in_bytes", @@ -89,6 +96,7 @@ __all__ = [ "WorkTileInfo", "PersistentTileSchedulerParams", "StaticPersistentTileScheduler", + "StaticPersistentRuntimeTileScheduler", "TensorMapUpdateMode", "TensorMapManager", "GroupSearchResult", @@ -96,6 +104,18 @@ __all__ = [ "create_initial_search_state", "GroupedGemmTileSchedulerHelper", "HardwareInfo", + "TransformMode", + "scale_tma_partition", + "transform_partition", + "scale_partition", + "get_gmem_layout_scale", + "get_smem_layout_scale", + "compute_smem_layout", + "get_transform_a_source", + "get_tma_atom_kind", + "get_copy_atom_a_transform", + "is_valid_scale_granularity", + "get_divisibility", "compute_epilogue_tile_shape", "get_smem_store_op", "get_tmem_load_op", @@ -107,4 +127,7 @@ __all__ = [ "make_blockscaled_trivial_tiled_mma", "sm90", "sm100", + "print_latex", + "print_latex_tv", + "distributed", ] diff --git a/python/CuTeDSL/cutlass/utils/blackwell_helpers.py b/python/CuTeDSL/cutlass/utils/blackwell_helpers.py index 13cac5c2..3c5c97b1 100644 --- a/python/CuTeDSL/cutlass/utils/blackwell_helpers.py +++ b/python/CuTeDSL/cutlass/utils/blackwell_helpers.py @@ -689,7 +689,7 @@ def make_smem_layout_a( is_k_major = tiled_mma.op.a_major_mode == OperandMajorMode.K a_smem_shape = tiled_mma.partition_shape_A( - cute.dice(mma_tiler_mnk, (1, None, 1), loc=loc, ip=ip) + cute.dice(mma_tiler_mnk, (1, None, 1), loc=loc, ip=ip), loc=loc, ip=ip ) a_smem_shape_mn_k = ( cute.size(a_smem_shape[0][0], loc=loc, ip=ip) * a_smem_shape[1], @@ -741,7 +741,7 @@ def make_smem_layout_b( is_k_major = tiled_mma.op.b_major_mode == OperandMajorMode.K b_smem_shape = tiled_mma.partition_shape_B( - cute.dice(mma_tiler_mnk, (None, 1, 1), loc=loc, ip=ip) + cute.dice(mma_tiler_mnk, (None, 1, 1), loc=loc, ip=ip), loc=loc, ip=ip ) b_smem_shape_nk = ( cute.size(b_smem_shape[0][0], loc=loc, ip=ip) * b_smem_shape[1], @@ -939,7 +939,9 @@ def make_trivial_tiled_mma( else: raise TypeError(f"unsupported ab_dtype, got {ab_dtype}") - return cute.make_tiled_mma(cute.make_mma_atom(mma_op), loc=loc, ip=ip) + return cute.make_tiled_mma( + cute.make_mma_atom(mma_op, loc=loc, ip=ip), loc=loc, ip=ip + ) @dsl_user_op @@ -1009,7 +1011,9 @@ def make_blockscaled_trivial_tiled_mma( else: raise TypeError(f"unsupported ab_dtype, got {ab_dtype}") - return cute.make_tiled_mma(cute.make_mma_atom(mma_op), loc=loc, ip=ip) + return cute.make_tiled_mma( + cute.make_mma_atom(mma_op, loc=loc, ip=ip), loc=loc, ip=ip + ) @dsl_user_op diff --git a/python/CuTeDSL/cutlass/utils/distributed.py b/python/CuTeDSL/cutlass/utils/distributed.py new file mode 100644 index 00000000..cd059c61 --- /dev/null +++ b/python/CuTeDSL/cutlass/utils/distributed.py @@ -0,0 +1,386 @@ +# 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 +# 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 functools import partial +from typing import Tuple + +import cutlass.cute as cute +from cutlass.cute.typing import Pointer, Int32 +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, nvvm +from cutlass._mlir.dialects.nvvm import ( + MemOrderKind, + MemScopeKind, + AtomicOpKind, +) + + +__all__ = [ + # misc + "atomicAdd", + "ld_bypass", + # Message Passing Lock & Unlock + "multimem_red_add1", + "red_add1", + "spin_lock_atom_cas_relaxed_wait", + # Load & Store + "multimem_ld_reduce_8xf16", + "multimem_ld_reduce_4xf32", + "multimem_ld_reduce_8xbf16", + "multimem_ld_reduce_16xe4m3", + "multimem_ld_reduce_16xe5m2", + "multimem_st_4xb32", +] + + +@dsl_user_op +def atomicAdd(dst_ptr: Pointer, val: Int32, loc=None, ip=None) -> Int32: + return nvvm.atomicrmw( + T.i32(), + AtomicOpKind.ADD, + dst_ptr.llvm_ptr, + val.ir_value(loc=loc, ip=ip), + mem_order=MemOrderKind.RELAXED, + syncscope=MemScopeKind.SYS, + loc=loc, + ip=ip, + ) + + +@cute.jit +def ld_bypass(input_tensor: cute.Tensor): + fragment = cute.make_rmem_tensor(input_tensor.layout, input_tensor.element_type) + copy_atom_load = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + input_tensor.element_type, + memory_order=cute.nvgpu.common.MemoryOrder.VOLATILE, + memory_scope=cute.nvgpu.common.MemoryScope.SYS, + ) + cute.copy_atom_call(copy_atom_load, input_tensor, fragment) + vals = fragment.load() + return vals + + +######################################################## +# Message Passing Lock & Unlock +######################################################## + + +@dsl_user_op +def multimem_red_release_gpu_add1( + lock_ptr: Pointer, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [lock_ptr.toint().ir_value(loc=loc, ip=ip)], + "multimem.red.release.gpu.global.add.u32 [$0], 1;", + "l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def multimem_red_release_sys_add1( + lock_ptr: Pointer, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [lock_ptr.toint().ir_value(loc=loc, ip=ip)], + "multimem.red.release.sys.global.add.u32 [$0], 1;", + "l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def multimem_red_relaxed_gpu_add1( + lock_ptr: Pointer, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [lock_ptr.toint().ir_value(loc=loc, ip=ip)], + "multimem.red.relaxed.gpu.global.add.u32 [$0], 1;", + "l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def multimem_red_relaxed_sys_add1( + lock_ptr: Pointer, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [lock_ptr.toint().ir_value(loc=loc, ip=ip)], + "multimem.red.relaxed.sys.global.add.u32 [$0], 1;", + "l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def multimem_red_add1( + lock_ptr: Pointer, + *, + order: str, + scope: str, + loc=None, + ip=None, +) -> None: + """ + add 1 to multicast ptr + """ + if scope == "gpu": + if order == "release": + multimem_red_release_gpu_add1(lock_ptr=lock_ptr, loc=loc, ip=ip) + elif order == "relaxed": + multimem_red_relaxed_gpu_add1(lock_ptr=lock_ptr, loc=loc, ip=ip) + elif scope == "sys": + if order == "release": + multimem_red_release_sys_add1(lock_ptr=lock_ptr, loc=loc, ip=ip) + elif order == "relaxed": + multimem_red_relaxed_sys_add1(lock_ptr=lock_ptr, loc=loc, ip=ip) + + +@dsl_user_op +def red_release_gpu_add1( + lock_ptr: Pointer, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [lock_ptr.toint().ir_value(loc=loc, ip=ip)], + "red.release.gpu.global.add.u32 [$0], 1;", + "l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_release_sys_add1( + lock_ptr: Pointer, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [lock_ptr.toint().ir_value(loc=loc, ip=ip)], + "red.release.sys.global.add.u32 [$0], 1;", + "l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_relaxed_gpu_add1( + lock_ptr: Pointer, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [lock_ptr.toint().ir_value(loc=loc, ip=ip)], + "red.relaxed.gpu.global.add.u32 [$0], 1;", + "l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_relaxed_sys_add1( + lock_ptr: Pointer, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [lock_ptr.toint().ir_value(loc=loc, ip=ip)], + "red.relaxed.sys.global.add.u32 [$0], 1;", + "l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_add1( + lock_ptr: Pointer, + *, + order: str, + scope: str, + loc=None, + ip=None, +) -> None: + """ + add 1 to multicast ptr + """ + if scope == "gpu": + if order == "release": + red_release_gpu_add1(lock_ptr=lock_ptr, loc=loc, ip=ip) + elif order == "relaxed": + red_relaxed_gpu_add1(lock_ptr=lock_ptr, loc=loc, ip=ip) + elif scope == "sys": + if order == "release": + red_release_sys_add1(lock_ptr=lock_ptr, loc=loc, ip=ip) + elif order == "relaxed": + red_relaxed_sys_add1(lock_ptr=lock_ptr, loc=loc, ip=ip) + + +@cute.jit +def spin_lock_atom_cas_relaxed_wait( + lock_ptr: Pointer, + *, + expected_val: Int32, + reset_val: Int32, + scope: str, + loc=None, + ip=None, +) -> None: + """ + wait on a spin lock until the expected count is reached. Reset flag to reset_val if the expected count is reached. + """ + if scope == "gpu": + result = 0 + while result != expected_val: + result = nvvm.atomicrmw( + T.i32(), + AtomicOpKind.CAS, + lock_ptr.llvm_ptr, + Int32(reset_val).ir_value(loc=loc, ip=ip), + b=Int32(expected_val).ir_value(loc=loc, ip=ip), + mem_order=MemOrderKind.RELAXED, + syncscope=MemScopeKind.GPU, + loc=loc, + ip=ip, + ) + elif scope == "sys": + result = 0 + while result != expected_val: + result = nvvm.atomicrmw( + T.i32(), + AtomicOpKind.CAS, + lock_ptr.llvm_ptr, + Int32(reset_val).ir_value(loc=loc, ip=ip), + b=Int32(expected_val).ir_value(loc=loc, ip=ip), + mem_order=MemOrderKind.RELAXED, + syncscope=MemScopeKind.SYS, + loc=loc, + ip=ip, + ) + + +######################################################## +# Multimem Load & Store +######################################################## + + +@dsl_user_op +def multimem_ld_reduce_base( + mc_ptr: Pointer, + *, + ptx_string: str = "", + loc=None, + ip=None, +) -> Tuple[Int32, Int32, Int32, Int32]: + # ld reduce 8xf16 elts + mc_ptr_int = mc_ptr.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) + return_struct = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(i32,i32,i32,i32)>"), + [mc_ptr_int], + ptx_string, + "=r,=r,=r,=r,l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return_regs = [llvm.extractvalue(T.i32(), return_struct, [i]) for i in range(4)] + return return_regs[0], return_regs[1], return_regs[2], return_regs[3] + + +multimem_ld_reduce_8xf16 = partial( + multimem_ld_reduce_base, + ptx_string="multimem.ld_reduce.sys.relaxed.global.add.acc::f32.v4.f16x2 {$0, $1, $2, $3}, [$4];", +) +multimem_ld_reduce_4xf32 = partial( + multimem_ld_reduce_base, + ptx_string="multimem.ld_reduce.sys.relaxed.global.add.v4.f32 {$0, $1, $2, $3}, [$4];", +) +multimem_ld_reduce_8xbf16 = partial( + multimem_ld_reduce_base, + ptx_string="multimem.ld_reduce.sys.relaxed.global.add.acc::f32.v4.bf16x2 {$0, $1, $2, $3}, [$4];", +) +multimem_ld_reduce_16xe4m3 = partial( + multimem_ld_reduce_base, + ptx_string="multimem.ld_reduce.sys.relaxed.global.add.acc::f16.v4.e4m3x4 {$0, $1, $2, $3}, [$4];", +) +multimem_ld_reduce_16xe5m2 = partial( + multimem_ld_reduce_base, + ptx_string="multimem.ld_reduce.sys.relaxed.global.add.acc::f16.v4.e5m2x4 {$0, $1, $2, $3}, [$4];", +) + + +@dsl_user_op +def multimem_st_4xb32( + mc_ptr: Pointer, + x: Int32, + y: Int32, + z: Int32, + w: Int32, + *, + loc=None, + ip=None, +) -> None: + # st 4x32 bits of data + mc_ptr_int = mc_ptr.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) + llvm.inline_asm( + T.i32(), + [mc_ptr_int, x, y, z, w], + "multimem.st.sys.relaxed.global.v4.f32 [$1], {$2, $3, $4, $5};", + "=r,l,r,r,r,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) diff --git a/python/CuTeDSL/cutlass/utils/grouped_gemm_tile_scheduler_helper.py b/python/CuTeDSL/cutlass/utils/grouped_gemm_tile_scheduler_helper.py index b990e960..615ccb3c 100644 --- a/python/CuTeDSL/cutlass/utils/grouped_gemm_tile_scheduler_helper.py +++ b/python/CuTeDSL/cutlass/utils/grouped_gemm_tile_scheduler_helper.py @@ -171,8 +171,18 @@ class GroupedGemmTileSchedulerHelper: def __new_from_mlir_values__( self, values: List[ir.Value] ) -> "GroupedGemmTileSchedulerHelper": - tile_sched_params = new_from_mlir_values(self.tile_sched_params, values) - search_state = new_from_mlir_values(self.search_state, values[1:]) + # Reconstruct tile_sched_params and determine how many values it consumed. + # NOTE: tile_sched_params may contain FastDivmod divisors (when swizzle_size == 1), + # which adds extra MLIR values. + params_values = extract_mlir_values(self.tile_sched_params) + n_params_values = len(params_values) + tile_sched_params = new_from_mlir_values( + self.tile_sched_params, values[:n_params_values] + ) + + # Reconstruct search_state from remaining values + search_state = new_from_mlir_values(self.search_state, values[n_params_values:]) + return GroupedGemmTileSchedulerHelper( self.group_count, tile_sched_params, diff --git a/python/CuTeDSL/cutlass/utils/hardware_info.py b/python/CuTeDSL/cutlass/utils/hardware_info.py index 2329b76d..a80d2471 100644 --- a/python/CuTeDSL/cutlass/utils/hardware_info.py +++ b/python/CuTeDSL/cutlass/utils/hardware_info.py @@ -9,7 +9,9 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -from cuda.bindings import driver, nvrtc +from cuda.bindings import driver, nvrtc, runtime +from cutlass.cutlass_dsl.cuda_jit_executor import CudaDialectJitModule +from cutlass.base_dsl.common import DSLRuntimeError import cutlass.cute as cute @@ -51,7 +53,7 @@ class HardwareInfo: f"Cluster size must be between 1 and 32, {cluster_size} is not supported" ) - device_fn = self._get_device_function(self.device) + self._get_device_function(self.device) max_shared_memory_per_block = self._checkCudaErrors( driver.cuDeviceGetAttribute( @@ -173,7 +175,14 @@ class HardwareInfo: ) # get a empty kernel to compute occupancy - def _get_device_function(self, device) -> None: + def _get_device_function(self, device) -> driver.CUfunction: self.compiled_kernel = cute.compile(self._host_function).to(device) - self.kernel = self.compiled_kernel.exec_context.kernel_functions[0] - self.module = self.compiled_kernel.exec_context.module.cuda_modules[0] + assert isinstance(self.compiled_kernel.jit_module, CudaDialectJitModule) + err, kernels = runtime.cudaLibraryEnumerateKernels( + 1, self.compiled_kernel.jit_module.cuda_library[0] + ) + if err is not runtime.cudaError_t.cudaSuccess: + raise DSLRuntimeError(f"Failed to enumerate kernels: {err}") + self.kernel = kernels[0] + self.kernel = self._checkCudaErrors(driver.cuKernelGetFunction(self.kernel)) + return self.kernel diff --git a/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py b/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py new file mode 100644 index 00000000..f40b9cdc --- /dev/null +++ b/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py @@ -0,0 +1,507 @@ +# 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. + +from __future__ import annotations + +from enum import Enum, auto +from math import log2 +from typing import Optional, Union + +import cutlass +import cutlass.cute as cute +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack + +""" +This file contain common utilities for mixed-input GEMM. +""" + + +class TransformMode(Enum): + """ + An enumeration for the possible transform modes of a mixed-input GEMM. + """ + + ConvertOnly = auto() + ConvertScale = auto() + + +def scale_tma_partition( + tCsS: cute.Tensor, + tCgS: cute.Tensor, + tma_atom_s: cute.CopyAtom, + block_in_cluster_coord_vmnk: cute.Coord, + scale_cta_layout: cute.Layout, +) -> tuple[cute.Tensor, cute.Tensor]: + """ + Perform TMA partition for scale tensor. + This method partitions the gobal memory and shared memory buffer for scale tensor for TMA load. + :param tCsS: Input scale shared memory tensor + :type tCsS: cute.Tensor + :param tCgS: Input scale global memory tensor + :type tCgS: cute.Tensor + :param tma_atom_s: TMA copy atom for scale tensor + :type tma_atom_s: cute.CopyAtom + :param block_in_cluster_coord_vmnk: CTA coord in the cluster + :type block_in_cluster_coord_vmnk: cute.Coord + :param scale_cta_layout: Layout of CTA from the view of the scale tensor + :type scale_cta_layout: cute.Layout + :return: A tuple containing (tSsS, tSgS) where: + + * tSsS: Partitioned scale tensor in shared memory + * tSgS: Partitioned scale tensor in global memory + + :rtype: tuple[cute.Tensor, cute.Tensor] + """ + tSsS, tSgS = cpasync.tma_partition( + tma_atom_s, + block_in_cluster_coord_vmnk[2], + scale_cta_layout, + cute.group_modes(tCsS, 0, 3), + cute.group_modes(tCgS, 0, 3), + ) + # Add rest_v mode + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tSsS = cute.make_tensor( + tSsS.iterator, + cute.make_layout( + ((tSsS.layout.shape[0], 1), *tSsS.layout.shape[1:]), + stride=( + (tSsS.layout.stride[0], 0), + *tSsS.layout.stride[1:], + ), + ), + ) + tSgS = cute.make_tensor( + tSgS.iterator, + cute.make_layout( + ((tSgS.layout.shape[0], 1), *tSgS.layout.shape[1:]), + stride=( + (tSgS.layout.stride[0], 0), + *tSgS.layout.stride[1:], + ), + ), + ) + return tSsS, tSgS + + +def transform_partition( + transform_a_source: tcgen05.OperandSource, + scale_mode: TransformMode, + copy_atom_a_input: cute.CopyAtom, + copy_atom_a_transform: cute.CopyAtom, + sA_input: cute.Tensor, + A_transform: cute.Tensor, + transform_local_tidx: cutlass.Int32, +) -> tuple[cute.TiledCopy, cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Partition tensors for transform input and output. + This method sets up the copy atoms and partitions the shared/tensor memory + for the transformation of tensor A. + :param transform_a_source: Where the transformed tensor A is stored (TMEM or SMEM) + :type transform_a_source: tcgen05.OperandSource + :param scale_mode: The transform mode (ConvertOnly or ConvertScale) + :type scale_mode: TransformMode + :param copy_atom_a_input: Copy atom for loading A from shared memory + :type copy_atom_a_input: cute.CopyAtom + :param copy_atom_a_transform: Copy atom for storing transformed A + :type copy_atom_a_transform: cute.CopyAtom + :param sA_input: Input tensor A in shared memory + :type sA_input: cute.Tensor + :param A_transform: Transformed tensor A in tensor or shared memory + :type A_transform: cute.Tensor + :param transform_local_tidx: Local thread index for transformation warps + :type transform_local_tidx: cutlass.Int32 + :return: A tuple containing (src_copy_a, dst_copy_a, tAsA_input, tA_transform) where: + + * src_copy_a: Tiled copy for source tensor + * dst_copy_a: Tiled copy for destination tensor + * tAsA_input: Partitioned input tensor A + * tA_transform: Partitioned transformed tensor A + + :rtype: tuple[cute.TiledCopy, cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + if cutlass.const_expr(transform_a_source == tcgen05.OperandSource.TMEM): + if cutlass.const_expr( + cute.size(A_transform, mode=[0, 0]) == 128 + and cute.size(sA_input, mode=[0, 0]) == 64 + ): + tensor_input = cute.make_tensor( + sA_input.iterator, + cute.logical_product( + sA_input.layout, + ((cute.make_layout(2, stride=0), None), None, None, None), + ), + ) + else: + tensor_input = sA_input + reg2tmem_tiled_copy = tcgen05.make_tmem_copy( + copy_atom_a_transform, A_transform[(None, None, None, 0)] + ) + thr_reg2tmem_tiled_copy = reg2tmem_tiled_copy.get_slice(transform_local_tidx) + partitioned_tensor_input = thr_reg2tmem_tiled_copy.partition_S(tensor_input) + partitioned_tensor_transform = thr_reg2tmem_tiled_copy.partition_D(A_transform) + src_copy_a = ( + cute.make_tiled_copy_S(copy_atom_a_input, reg2tmem_tiled_copy) + if scale_mode is TransformMode.ConvertScale + else None + ) + dst_copy_a = reg2tmem_tiled_copy + tAsA_input = partitioned_tensor_input + tA_transform = partitioned_tensor_transform + elif cutlass.const_expr(transform_a_source == tcgen05.OperandSource.SMEM): + # Construct tiled_copy satisfying 8 contiguous elts per copy atom + reg2smem_tiled_copy = cute.make_cotiled_copy( + copy_atom_a_transform, + cute.make_layout((128, 8), stride=(8, 1)), + A_transform[(None, None, None, 0)].layout, + ) + thr_reg2smem_tiled_copy = reg2smem_tiled_copy.get_slice(transform_local_tidx) + partitioned_tensor_input = thr_reg2smem_tiled_copy.partition_S(sA_input) + partitioned_tensor_transform = thr_reg2smem_tiled_copy.partition_D(A_transform) + src_copy_a = ( + cute.make_tiled_copy_S(copy_atom_a_input, reg2smem_tiled_copy) + if scale_mode is TransformMode.ConvertScale + else None + ) + # auto-vec copy is enough for copy from register to shared memory here + dst_copy_a = None + tAsA_input = partitioned_tensor_input + tA_transform = partitioned_tensor_transform + return src_copy_a, dst_copy_a, tAsA_input, tA_transform + + +def scale_partition( + src_copy_a: cute.TiledCopy, + tCsS: cute.Tensor, + transform_local_tidx: cutlass.Int32, + mma_dtype: type[cutlass.Numeric], +) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor]: + """ + Partition the scale tensor for transformation. + This method prepares the copy atom and partitions the shared memory for the scale tensor. + :param src_copy_a: Tiled copy for the source tensor + :type src_copy_a: cute.TiledCopy + :param tCsS: Scale tensor in shared memory + :type tCsS: cute.Tensor + :param transform_local_tidx: Local thread index for transformation warps + :type transform_local_tidx: cutlass.Int32 + :param mma_dtype: Data type for the MMA operation + :type mma_dtype: type[cutlass.Numeric] + :return: A tuple containing (smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS) where: + + * smem_thr_copy_S: Tiled copy for the scale tensor + * tSsS_trans: Partitioned scale tensor for transformation + * tSrS_copy: Register fragment for the scale tensor + * tSrS: View of scale tensor used for transformation computation + + :rtype: tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor] + """ + smem_thr_copy_S = None + tSsS_trans = None + tSrS = None + # Partition scale tensor + smem_thr_copy_S = src_copy_a.get_slice(transform_local_tidx) + tSsS_trans = smem_thr_copy_S.partition_S(tCsS) + # Construct register fragment for scale tensor + tSsS_layout_per_stage = tSsS_trans[(None, None, None, None, 0)].layout + # tSrS for copy + tSrS_copy = cute.make_rmem_tensor( + cute.filter_zeros(tSsS_layout_per_stage).shape, mma_dtype + ) + # tSrS view for transformation computation + tSrS = cute.make_tensor( + tSrS_copy.iterator, + cute.make_layout(tSsS_layout_per_stage.shape, stride=tSrS_copy.layout.stride), + ) + return smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS + + +def get_gmem_layout_scale( + scale_shape_mkl: tuple[int, int, int], + scale_granularity_m: int, + scale_granularity_k: int, + scale_major_mode: tcgen05.OperandMajorMode, +) -> cute.Layout: + """ + Get the layout of the scale tensor in global memory. + :param scale_shape_mkl: The shape of the scale tensor (M, K, L). + :type scale_shape_mkl: tuple[int, int, int] + :return: The layout of the scale tensor in global memory. + :rtype: cute.Layout + """ + m, k, l = scale_shape_mkl + shape_scale = ( + (scale_granularity_m, cute.ceil_div(m, scale_granularity_m)), + (scale_granularity_k, cute.ceil_div(k, scale_granularity_k)), + ) + if cutlass.const_expr(scale_major_mode == tcgen05.OperandMajorMode.MN): + layout_mk = cute.make_layout( + shape_scale, + stride=( + (0, 1), + (0, cute.size(shape_scale[0][1])), + ), + ) + else: + layout_mk = cute.make_layout( + shape_scale, + stride=( + (0, cute.size(shape_scale[1][1])), + (0, 1), + ), + ) + return cute.make_layout( + (*layout_mk.shape, l), + stride=(*layout_mk.stride, cute.cosize(layout_mk)), + ) + + +def get_smem_layout_scale( + mma_tiler: tuple[int, int, int], + use_2cta_instrs: bool, + scale_granularity_m: int, + scale_granularity_k: int, + scale_major_mode: tcgen05.OperandMajorMode, + a_scale_dtype: type[cutlass.Numeric], + num_scale_load2trans_stage: int, +) -> tuple[tuple[int, int], cute.ComposedLayout, cute.ComposedLayout]: + """ + Get the layout of the scale tensor in shared memory. + :return: A tuple containing (scale_tile_shape, smem_layout_scale_per_stage, smem_layout_scale) where: + + * scale_tile_shape: The tile shape + * smem_layout_scale_per_stage: Shared memory layout for scale tensor per stage + * smem_layout_scale: Shared memory layout for scale tensor + + :rtype: tuple[tuple[int, int], cute.ComposedLayout, cute.ComposedLayout] + """ + scale_tile_shape = ( + (cute.size(mma_tiler[0]) // 2 if use_2cta_instrs else cute.size(mma_tiler[0])), + cute.size(mma_tiler[2]), + ) + size_mn = scale_tile_shape[0] + size_k = scale_tile_shape[1] + smem_size_mn = scale_granularity_m if scale_granularity_m < size_mn else size_mn + smem_size_k = scale_granularity_k if scale_granularity_k < size_k else size_k + div_mn = cute.ceil_div(size_mn, smem_size_mn) + div_k = cute.ceil_div(size_k, smem_size_k) + smem_atom_shape = ( + (smem_size_mn, div_mn), + (smem_size_k, div_k), + ) + if cutlass.const_expr(scale_major_mode == tcgen05.OperandMajorMode.MN): + outer_layout = cute.make_layout( + smem_atom_shape, + stride=( + (0, 1), + (0, div_mn), + ), + ) + else: + outer_layout = cute.make_layout( + smem_atom_shape, + stride=( + (0, div_k), + (0, 1), + ), + ) + # Apply a trivial swizzle to make it a composed layout, which could be used to construct TMA atom + smem_layout_scale_per_stage = cute.make_composed_layout( + cute.make_swizzle(0, 4, 3), 0, outer_layout + ) + assert cute.rank(smem_layout_scale_per_stage) == 2, "Scale layout must be rank 2" + assert ( + cute.size(mma_tiler[0]) % cute.size(smem_layout_scale_per_stage.outer[0]) == 0 + ), "smem_layout_scale_per_stage must equal the tile shape." + assert ( + 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" + ) + # Scale layout in smem with multiple stages + smem_layout_scale = cute.append( + smem_layout_scale_per_stage, + cute.make_layout( + (num_scale_load2trans_stage), + stride=(cute.cosize(smem_layout_scale_per_stage.outer)), + ), + ) + return scale_tile_shape, smem_layout_scale_per_stage, smem_layout_scale + + +def compute_smem_layout( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: tuple[int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + load2trans_stage_count: int, + trans2mma_stage_count: int, +) -> tuple[ + cute.ComposedLayout, + cute.ComposedLayout, + cute.ComposedLayout, +]: + """ + Compute shared memory layouts for tensor A, transformed A and tensor B. + :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 load2trans_stage_count: Number of stages for load-to-transform pipeline. + :type load2trans_stage_count: int + :param trans2mma_stage_count: Number of stages for transform-to-MMA pipeline. + :type trans2mma_stage_count: int + :return: A tuple containing (smem_layout_a, smem_layout_a_transform, smem_layout_b) where: + + * smem_layout_a: Shared memory layout for tensor A + * smem_layout_a_transform: Shared memory layout for transformed tensor A + * smem_layout_b: Shared memory layout for tensor B + + :rtype: tuple[cute.ComposedLayout, cute.ComposedLayout, cute.ComposedLayout] + """ + smem_layout_a = sm100_utils.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + a_dtype, + load2trans_stage_count, + ) + smem_layout_a_transform = sm100_utils.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + tiled_mma.op.a_dtype, + trans2mma_stage_count, + ) + smem_layout_b = sm100_utils.make_smem_layout_b( + tiled_mma, + mma_tiler_mnk, + b_dtype, + load2trans_stage_count, + ) + return (smem_layout_a, smem_layout_a_transform, smem_layout_b) + + +def get_transform_a_source( + a_major_mode: tcgen05.OperandMajorMode, +) -> tcgen05.OperandSource: + """ + Determine the operand source for transformed A tensor based on the operand major mode. + """ + if cutlass.const_expr(a_major_mode == tcgen05.OperandMajorMode.K): + return tcgen05.OperandSource.TMEM + else: + return tcgen05.OperandSource.SMEM + + +def get_tma_atom_kind( + mcast: cutlass.Boolean, + use_2cta_instrs: bool, + is_b: bool, +) -> Union[cpasync.CopyBulkTensorTileG2SMulticastOp, cpasync.CopyBulkTensorTileG2SOp]: + """ + Get the TMA atom kind based on 1) whether it's a multicast operation, + 2) whether 2CTA tcgen05.mma instruction is enabled, and + 3) whether it's a B tensor + """ + # Not using .2CTA instructions for tensor A as the consumer is threads on different CTAs + cta_group = ( + tcgen05.CtaGroup.TWO if (use_2cta_instrs and is_b) else tcgen05.CtaGroup.ONE + ) + if cutlass.const_expr(mcast): + return cpasync.CopyBulkTensorTileG2SMulticastOp(cta_group) + return cpasync.CopyBulkTensorTileG2SOp(cta_group) + + +def get_copy_atom_a_transform( + mma_dtype: type[cutlass.Numeric], + use_2cta_instrs: bool, + transform_a_source: tcgen05.OperandSource, + a_smem_shape: cute.Shape, + a_dtype: type[cutlass.Numeric], +) -> cute.CopyAtom: + """ + Determine the copy atom for transformed A tensor based on the operand source and tile size. + """ + if cutlass.const_expr(transform_a_source == tcgen05.OperandSource.TMEM): + if cutlass.const_expr( + cute.size(a_smem_shape[0][0]) == 64 and (not use_2cta_instrs) + ): + copy_op_r2t = tcgen05.St16x256bOp( + tcgen05.Repetition(1), tcgen05.Unpack.NONE + ) + else: + copy_op_r2t = tcgen05.St32x32bOp(tcgen05.Repetition(8), tcgen05.Unpack.NONE) + return cute.make_copy_atom(copy_op_r2t, mma_dtype) + else: + return cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), a_dtype, num_bits_per_copy=32 + ) + + +def is_valid_scale_granularity( + scale_granularity_m: int, + scale_granularity_k: int, + a_dtype: type[cutlass.Numeric], + k: int, + mma_tiler_k: int, +) -> bool: + """ + Check if the scale granularity settings are valid for the given data type and problem size. + """ + if a_dtype.width == 8: + # No scale tensor for 8bit data type A + if not (scale_granularity_m == 0 and scale_granularity_k == 0): + return False + elif a_dtype.width == 4: + if scale_granularity_m != 1 or ( + scale_granularity_k == 0 + or k % scale_granularity_k != 0 + or scale_granularity_k % mma_tiler_k != 0 + ): + return False + return True + + +def get_divisibility(contiguous_dim_size: int, upper_bound: int = 128) -> int: + """ + Calculate the largest power of 2 divisibility factor for memory alignment. + """ + # Check the largest power of 2 factor of contiguous_dim_size + for i in range(int(log2(contiguous_dim_size)), 0, -1): + if contiguous_dim_size % (2**i) == 0: + return min(2**i, upper_bound) + return 1 diff --git a/python/CuTeDSL/cutlass/utils/print_latex.py b/python/CuTeDSL/cutlass/utils/print_latex.py new file mode 100644 index 00000000..817b7bea --- /dev/null +++ b/python/CuTeDSL/cutlass/utils/print_latex.py @@ -0,0 +1,165 @@ +# 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 +# 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 Callable, Union + +from ..cute import ( + Layout, + ComposedLayout, + append, + is_static, + make_layout, + size, + product_each, + rank, +) +from ..cute.typing import IntTuple + +__all__ = ["print_latex", "print_latex_tv"] + + +def tikz_color_bwx8(idx: int): + color_map = [ + "black!00", + "black!40", + "black!20", + "black!60", + "black!10", + "black!50", + "black!30", + "black!70", + ] + return color_map[idx % 8] + + +def tikz_color_white(idx: int): + return "white" + + +def tikz_color_tv(tid: int, vid: int): + color_map = [ + "{rgb,255:red,175;green,175;blue,255}", + "{rgb,255:red,175;green,255;blue,175}", + "{rgb,255:red,255;green,255;blue,175}", + "{rgb,255:red,255;green,175;blue,175}", + "{rgb,255:red,210;green,210;blue,255}", + "{rgb,255:red,210;green,255;blue,210}", + "{rgb,255:red,255;green,255;blue,210}", + "{rgb,255:red,255;green,210;blue,210}", + ] + return color_map[tid % 8] + + +def print_latex(x: Union[Layout, ComposedLayout], *, color: Callable = tikz_color_bwx8): + """ + Prints a layout. + :param x: A layout + :type x: Union[Layout, ComposedLayout] + :param color: A function that returns TiKZ colors + :type color: Callable + """ + + if not is_static(x): + raise ValueError("Requires static input") + if rank(x) > 2: + raise ValueError("Requires rank <= 2 to print") + + if rank(x) == 1: + layout = append(x, make_layout(1, stride=0)) + else: + layout = x + + print("%% Layout: {}", layout) + print("\\documentclass[convert]{standalone}") + print("\\usepackage{tikz}") + print("\\begin{document}") + print( + "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]" + ) + + M, N = product_each(x.shape) + + for m in range(M): + for n in range(N): + idx = layout((m, n)) + print("\\node[fill=") + print(color(idx)) + print("] at (%d,%d) {%d};\n" % (m, n, idx)) + print( + "\\draw[color=black,thick,shift={(-0.5,-0.5)}] (0,0) grid (%d,%d);\n\n" % (M, N) + ) + for m in range(M): + print("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (m, -1, m)) + for n in range(N): + print("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (-1, n, n)) + + ## Footer + print("\\end{tikzpicture}") + print("\\end{document}") + + +def print_latex_tv( + layout_tv: Union[Layout, ComposedLayout], + tile_mn: Union[IntTuple, Layout], + *, + color: Callable = tikz_color_tv, +): + """ + Prints a tv layout for a tile M N. Everything must be static. + :param layout_tv: A static thread value layout + :type layout_tv: Union[Layout, ComposedLayout] + :param tile_mn: A static M N tile + :type tile_mn: Union[IntTuple, Layout] + :param color: A function that returns TiKZ colors + :type color: Callable + """ + if not is_static(layout_tv) or not is_static(tile_mn): + raise ValueError("Layout tv and tile_mn must be static") + if rank(layout_tv) != 2: + raise ValueError("Require layout_tv to be rank 2") + + print("%% Layout TV: {}", layout_tv) + print("\\documentclass[convert]{standalone}") + print("\\usepackage{tikz}") + print("\\begin{document}") + print( + "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]\n" + ) + + if not isinstance(tile_mn, Layout): + tile_mn = make_layout(tile_mn) + + M, N = product_each(tile_mn.shape) + filled = [[False for n in range(N)] for m in range(M)] + + for tid in range(size(layout_tv, mode=[0])): + for vid in range(size(layout_tv, mode=[1])): + idx = layout_tv((tid, vid)) + m = (idx // tile_mn.stride[0]) % tile_mn.shape[0] + n = (idx // tile_mn.stride[1]) % tile_mn.shape[1] + if not filled[m][n]: + filled[m][n] = True + print( + "\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n" + % (color(tid, vid), m, n, tid, vid) + ) + + print( + "\\draw[color=black,thick,shift={(-0.5,-0.5)}] (0,0) grid (%d,%d);\n\n" % (M, N) + ) + for m in range(M): + print("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (m, -1, m)) + for n in range(N): + print("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (-1, n, n)) + + ## Footer + print("\\end{tikzpicture}") + print("\\end{document}") diff --git a/python/CuTeDSL/cutlass/utils/smem_allocator.py b/python/CuTeDSL/cutlass/utils/smem_allocator.py index bd62ed0e..7e801ddd 100644 --- a/python/CuTeDSL/cutlass/utils/smem_allocator.py +++ b/python/CuTeDSL/cutlass/utils/smem_allocator.py @@ -10,6 +10,7 @@ # is strictly prohibited. from typing import Optional, Type, Union, overload +import inspect import cutlass.cute as cute from cutlass.cute.arch import get_dyn_smem, get_dyn_smem_size @@ -285,4 +286,11 @@ class SmemAllocator: return cute.make_tensor(ptr, layout, loc=loc, ip=ip) +# Set explicit signature for Sphinx documentation to avoid issues with @dsl_user_op decorator +SmemAllocator.__init__.__signature__ = inspect.Signature( + [ + inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD), + ] +) + get_smem_capacity_in_bytes = SmemAllocator.capacity_in_bytes diff --git a/python/CuTeDSL/cutlass/utils/static_persistent_tile_scheduler.py b/python/CuTeDSL/cutlass/utils/static_persistent_tile_scheduler.py index e42f4d2f..7b59575e 100644 --- a/python/CuTeDSL/cutlass/utils/static_persistent_tile_scheduler.py +++ b/python/CuTeDSL/cutlass/utils/static_persistent_tile_scheduler.py @@ -19,6 +19,7 @@ from cutlass.cutlass_dsl import ( extract_mlir_values, new_from_mlir_values, dsl_user_op, + const_expr, ) from cutlass._mlir import ir import cutlass.cute as cute @@ -135,6 +136,7 @@ class PersistentTileSchedulerParams: ip=ip, ) + # Apply swizzle if swizzle_size > 1 if swizzle_size > 1: problem_shape_ncluster_mnl = cute.round_up( self.problem_layout_ncluster_mnl.shape, @@ -172,6 +174,35 @@ class PersistentTileSchedulerParams: ip=ip, ) + # Create FastDivmod divisors (only when swizzle_size == 1 for correctness) + # FastDivmod assumes simple col-major layout, incompatible with swizzled layouts + if swizzle_size == 1: + problem_layout_size = cute.size( + self.problem_layout_ncluster_mnl, loc=loc, ip=ip + ) + cluster_count_m = self.problem_layout_ncluster_mnl.shape[0] + cluster_count_n = self.problem_layout_ncluster_mnl.shape[1] + + # batch_fdd: Used to map linear_idx to work_unit_id (handles persistent scheduling) + self.batch_fdd = cute.fast_divmod_create_divisor( + problem_layout_size, loc=loc, ip=ip + ) + + # cluster_shape_m_fdd: Used to decode work_unit_id to cluster coordinates + self.cluster_shape_m_fdd = cute.fast_divmod_create_divisor( + cluster_count_m, loc=loc, ip=ip + ) + + # cluster_shape_n_fdd: Used for the second level decomposition + self.cluster_shape_n_fdd = cute.fast_divmod_create_divisor( + cluster_count_n, loc=loc, ip=ip + ) + else: + # FastDivmod not applicable with swizzling, set to None + self.batch_fdd = None + self.cluster_shape_m_fdd = None + self.cluster_shape_n_fdd = None + def __extract_mlir_values__(self): values, self._values_pos = [], [] for obj in [ @@ -183,10 +214,38 @@ class PersistentTileSchedulerParams: obj_values = extract_mlir_values(obj) values += obj_values self._values_pos.append(len(obj_values)) + + # Add FastDivmod divisors to MLIR values for Host->Device transfer + # Only add non-None values to avoid MLIR type errors + fastdivmod_values = [] + fastdivmod_indices = [] # Track which FastDivmod objects are present + + for i, (fdd_name, fdd_obj) in enumerate( + [ + ("batch_fdd", self.batch_fdd), + ("cluster_shape_m_fdd", self.cluster_shape_m_fdd), + ("cluster_shape_n_fdd", self.cluster_shape_n_fdd), + ] + ): + if fdd_obj is not None: + # Extract MLIR values from FastDivmodDivisor objects + fdd_values = extract_mlir_values(fdd_obj) + fastdivmod_values.extend(fdd_values) + fastdivmod_indices.append(i) + + values += fastdivmod_values + self._values_pos.append( + len(fastdivmod_indices) + ) # Store count of FastDivmod objects, not values + self._fastdivmod_indices = fastdivmod_indices # Store for reconstruction + return values def __new_from_mlir_values__(self, values): obj_list = [] + values_copy = list(values) # Make a copy to avoid modifying original + + # Reconstruct original objects from MLIR values for obj, n_items in zip( [ self.problem_shape_ntile_mnl, @@ -194,11 +253,32 @@ class PersistentTileSchedulerParams: self.swizzle_size, self._raster_along_m, ], - self._values_pos, + self._values_pos[:-1], # Exclude FastDivmod count ): - obj_list.append(new_from_mlir_values(obj, values[:n_items])) - values = values[n_items:] - return PersistentTileSchedulerParams(*(tuple(obj_list)), loc=self._loc) + obj_list.append(new_from_mlir_values(obj, values_copy[:n_items])) + values_copy = values_copy[n_items:] + + # Create new params object by calling __init__ with reconstructed values + # This properly recreates layouts and other derived attributes in the device context + new_params = PersistentTileSchedulerParams(*(tuple(obj_list)), loc=self._loc) + + # Restore FastDivmod divisors from remaining values + fdd_names = ["batch_fdd", "cluster_shape_m_fdd", "cluster_shape_n_fdd"] + + if hasattr(self, "_fastdivmod_indices") and len(self._fastdivmod_indices) > 0: + # Override the FastDivmod divisors created by __init__ with reconstructed ones + for j, original_index in enumerate(self._fastdivmod_indices): + fdd_name = fdd_names[original_index] + # Get the original FastDivmodDivisor object + original_fdd = getattr(self, fdd_name) + if original_fdd is not None and j < len(values_copy): + # Each FastDivmodDivisor has 1 MLIR value + reconstructed_fdd = new_from_mlir_values( + original_fdd, [values_copy[j]] + ) + setattr(new_params, fdd_name, reconstructed_fdd) + + return new_params @dsl_user_op def get_grid_shape( @@ -285,12 +365,16 @@ class StaticPersistentTileScheduler: values.extend(extract_mlir_values(self._current_work_linear_idx)) values.extend(extract_mlir_values(self.cta_id_in_cluster)) values.extend(extract_mlir_values(self._num_tiles_executed)) + + # CRITICAL: Also extract FastDivmod divisors from params + values.extend(extract_mlir_values(self.params)) + return values def __new_from_mlir_values__( self, values: list[ir.Value] ) -> "StaticPersistentTileScheduler": - assert len(values) == 6 + assert len(values) >= 6 new_num_persistent_clusters = new_from_mlir_values( self.num_persistent_clusters, [values[0]] ) @@ -303,8 +387,13 @@ class StaticPersistentTileScheduler: new_num_tiles_executed = new_from_mlir_values( self._num_tiles_executed, [values[5]] ) + + # Reconstruct params with FastDivmod divisors + params_values = values[6:] # Remaining values are from params + new_params = new_from_mlir_values(self.params, params_values) + return StaticPersistentTileScheduler( - self.params, + new_params, # Use reconstructed params with FastDivmod divisors new_num_persistent_clusters, new_current_work_linear_idx, new_cta_id_in_cluster, @@ -403,11 +492,14 @@ class StaticPersistentTileScheduler: self.params.problem_layout_ncluster_mnl, loc=loc, ip=ip ) + # Choose coordinate calculation method based on swizzle configuration if self.params.swizzle_size == 1: - cur_cluster_coord = self.params.problem_layout_ncluster_mnl.get_hier_coord( + # Use FastDivmod optimization for non-swizzled layouts + cur_cluster_coord = self._get_cluster_work_idx_with_fastdivmod( current_work_linear_idx, loc=loc, ip=ip ) else: + # Use get_flat_coord for swizzled layouts (FastDivmod doesn't support them) cur_cluster_coord = self.params.problem_layout_ncluster_mnl.get_flat_coord( current_work_linear_idx, loc=loc, ip=ip ) @@ -424,6 +516,40 @@ class StaticPersistentTileScheduler: return WorkTileInfo(cur_tile_coord, is_valid) + def _get_cluster_work_idx_with_fastdivmod( + self, current_work_linear_idx: Int32, *, loc=None, ip=None + ) -> Tuple[Int32, Int32, Int32]: + """ + FastDivmod optimized CLUSTER coordinate calculation. + + CRITICAL: This should mimic problem_layout_ncluster_mnl.get_hier_coord() + which returns CLUSTER coordinates, not tile coordinates! + + :param current_work_linear_idx: Linear index in the work space + :type current_work_linear_idx: Int32 + :return: Cluster coordinates (m, n, l) or None if FastDivmod not available + :rtype: Tuple[Int32, Int32, Int32] or None + """ + + # Step 1: Handle persistent scheduling - map linear_idx to work_unit_id + work_iteration, work_unit_id = divmod( + current_work_linear_idx, self.params.batch_fdd + ) + + # Step 2: Decode work_unit_id using FastDivmod objects + # The layout structure is: problem_layout_ncluster_mnl has shape (cluster_count_m, cluster_count_n, batch_count) + # work_unit_id needs to be decomposed into (batch_l, cluster_n, cluster_m) in little-endian order + + # First, get cluster_m using cluster_shape_m_fdd + cluster_n_batch, cluster_m = divmod( + work_unit_id, self.params.cluster_shape_m_fdd + ) + + # Then decode cluster_n_batch to get cluster_n and batch_l using FastDivmod + batch_l, cluster_n = divmod(cluster_n_batch, self.params.cluster_shape_n_fdd) + + return (cluster_m, cluster_n, batch_l) + @dsl_user_op def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: return self._get_current_work_for_linear_idx( @@ -446,3 +572,176 @@ class StaticPersistentTileScheduler: return self._num_tiles_executed +class StaticPersistentRuntimeTileScheduler(StaticPersistentTileScheduler): + """A scheduler for static persistent runtime tile execution in CUTLASS/CuTe kernels. + This scheduler will always launch all the SMs and the scheduler will generate the real tile info for each SM. + + :ivar params: Tile schedule related params, including cluster shape and problem_layout_ncluster_mnl + :type params: PersistentTileSchedulerParams + :ivar num_persistent_clusters: Number of persistent clusters that can be launched + :type num_persistent_clusters: Int32 + :ivar cta_id_in_cluster: ID of the CTA within its cluster + :type cta_id_in_cluster: cute.Coord + :ivar _num_tiles_executed: Counter for executed tiles + :type _num_tiles_executed: Int32 + :ivar _current_work_linear_idx: Current cluster index + :type _current_work_linear_idx: Int32 + """ + + def __init__( + self, + params: PersistentTileSchedulerParams, + num_persistent_clusters: Int32, + current_work_linear_idx: Int32, + cta_id_in_cluster: cute.Coord, + num_tiles_executed: Int32, + inner_mode: int = 1, + ): + """ + Initializes the StaticPersistentTileScheduler with the given parameters. + + :param params: Tile schedule related params, including cluster shape and problem_layout_ncluster_mnl. + :type params: PersistentTileSchedulerParams + :param num_persistent_clusters: Number of persistent clusters that can be launched. + :type num_persistent_clusters: Int32 + :param current_work_linear_idx: Current cluster index. + :type current_work_linear_idx: Int32 + :param cta_id_in_cluster: ID of the CTA within its cluster. + :type cta_id_in_cluster: cute.Coord + :param num_tiles_executed: Counter for executed tiles. + :type num_tiles_executed: Int32 + :param inner_mode: The inner mode along which the linear index will be decomposed first. + :type inner_mode: int + """ + super().__init__( + params, + num_persistent_clusters, + current_work_linear_idx, + cta_id_in_cluster, + num_tiles_executed, + ) + if inner_mode not in [0, 1]: + raise ValueError( + f"inner_mode must be 0(for M mode) or 1(for N mode), but got {inner_mode}" + ) + self.inner_mode = inner_mode + + def __new_from_mlir_values__( + self, values: list[ir.Value] + ) -> "StaticPersistentRuntimeTileScheduler": + assert len(values) >= 6 + new_num_persistent_clusters = new_from_mlir_values( + self.num_persistent_clusters, [values[0]] + ) + new_current_work_linear_idx = new_from_mlir_values( + self._current_work_linear_idx, [values[1]] + ) + new_cta_id_in_cluster = new_from_mlir_values( + self.cta_id_in_cluster, values[2:5] + ) + new_num_tiles_executed = new_from_mlir_values( + self._num_tiles_executed, [values[5]] + ) + + # Reconstruct params with FastDivmod divisors (same as parent class) + params_values = values[6:] # Remaining values are from params + new_params = new_from_mlir_values(self.params, params_values) + + return StaticPersistentRuntimeTileScheduler( + new_params, # Use reconstructed params with FastDivmod divisors + new_num_persistent_clusters, + new_current_work_linear_idx, + new_cta_id_in_cluster, + new_num_tiles_executed, + self.inner_mode, + ) + + @staticmethod + @dsl_user_op + def create( + params: PersistentTileSchedulerParams, + block_idx: Tuple[Integer, Integer, Integer], + grid_dim: Tuple[Integer, Integer, Integer], + inner_mode: int = 1, + *, + loc=None, + ip=None, + ): + """Initialize the static persistent tile scheduler. + + :param params: Parameters for the persistent + tile scheduler. + :type params: PersistentTileSchedulerParams + :param block_idx: The 3d block index in the format (bidx, bidy, bidz). + :type block_idx: Tuple[Integer, Integer, Integer] + :param grid_dim: The 3d grid dimensions for kernel launch. + :type grid_dim: Tuple[Integer, Integer, Integer] + :param inner_mode: The inner mode along which the linear index will be decomposed first. + :type inner_mode: int + + :return: A StaticPersistentRuntimeTileScheduler object. + :rtype: StaticPersistentRuntimeTileScheduler + """ + + # Calculate the number of persistent clusters by dividing the total grid size + # by the number of CTAs per cluster + num_persistent_clusters = cute.size(grid_dim, loc=loc, ip=ip) // cute.size( + params.cluster_shape_mn, loc=loc, ip=ip + ) + + bidx, bidy, bidz = block_idx + + # Initialize workload index equals to the cluster index in the grid + current_work_linear_idx = Int32(bidz) + + # CTA id in the cluster + cta_id_in_cluster = ( + Int32(bidx % params.cluster_shape_mn[0]), + Int32(bidy % params.cluster_shape_mn[1]), + Int32(0), + ) + # Initialize number of tiles executed to zero + num_tiles_executed = Int32(0) + return StaticPersistentRuntimeTileScheduler( + params, + num_persistent_clusters, + current_work_linear_idx, + cta_id_in_cluster, + num_tiles_executed, + inner_mode, + ) + + # private method + def _get_current_work_for_linear_idx( + self, current_work_linear_idx: Int32, *, loc=None, ip=None + ) -> WorkTileInfo: + """Compute current tile coord given current_work_linear_idx and cta_id_in_cluster. + + :param current_work_linear_idx: The linear index of the current work. + :type current_work_linear_idx: Int32 + + :return: An object containing information about the current tile coordinates + and validity status. + :rtype: WorkTileInfo + """ + ntile_shape = self.params.problem_layout_ncluster_mnl.shape + int_max = 2147483647 + if const_expr(self.inner_mode == 1): + ntile_layout = cute.make_layout( + (int_max, ntile_shape[1]), stride=(ntile_shape[1], 1) + ) + else: + ntile_layout = cute.make_layout( + (ntile_shape[0], int_max), stride=(1, ntile_shape[0]) + ) + cluster_tile_coord_mn = ntile_layout.get_hier_coord(current_work_linear_idx) + cur_tile_coord = ( + cluster_tile_coord_mn[0], + cluster_tile_coord_mn[1], + Int32(0), + ) + + # it is determined by kernel implementation + is_valid = True + + return WorkTileInfo(cur_tile_coord, is_valid) diff --git a/python/CuTeDSL/cutlass/utils/tmem_allocator.py b/python/CuTeDSL/cutlass/utils/tmem_allocator.py index 6b832436..47b4338d 100644 --- a/python/CuTeDSL/cutlass/utils/tmem_allocator.py +++ b/python/CuTeDSL/cutlass/utils/tmem_allocator.py @@ -10,6 +10,7 @@ # is strictly prohibited. from typing import Optional, Type +import inspect from cutlass import const_expr from cutlass.cutlass_dsl import ( @@ -17,6 +18,7 @@ from cutlass.cutlass_dsl import ( Float32, extract_mlir_values, new_from_mlir_values, + dsl_user_op, ) import cutlass.pipeline as pipeline import cutlass.cute as cute @@ -43,22 +45,27 @@ class TmemAllocator: :type _two_cta_tmem_dealloc_mbar_ptr: cute.Pointer """ + @dsl_user_op @cute.jit - def _init_dealloc_mbarrier(self): + def _init_dealloc_mbarrier(self, *, loc=None, ip=None): assert self._two_cta_tmem_dealloc_mbar_ptr is not None, ( "two_cta_tmem_dealloc_mbar_ptr is required for two cta" ) - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) + warp_idx = cute.arch.warp_idx(loc=loc, ip=ip) + warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip) _is_allocator_warp = warp_idx == self._allocator_warp_id if _is_allocator_warp: num_tmem_dealloc_threads = 32 - with cute.arch.elect_one(): + with cute.arch.elect_one(loc=loc, ip=ip): cute.arch.mbarrier_init( - self._two_cta_tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads + self._two_cta_tmem_dealloc_mbar_ptr, + num_tmem_dealloc_threads, + loc=loc, + ip=ip, ) cute.arch.mbarrier_init_fence() + @dsl_user_op def __init__( self, alloc_result_dst_smem_ptr: cute.Pointer, @@ -67,11 +74,40 @@ class TmemAllocator: is_two_cta: bool = False, num_allocated_columns: int = 0, two_cta_tmem_dealloc_mbar_ptr: Optional[cute.Pointer] = None, + *, + loc=None, + ip=None, ): - """Initialize the TmemAllocator instance. + """ + Initialize a TmemAllocator instance for managing tensor memory on Blackwell GPUs. - Sets up the allocator state by initializing smem pointer that holds the base address of allocated tensor memory, allocator warp id, whether it is for two cta, number of allocated columns, and barrier for retrieving tensor memory ptr. - Meanwhile, it also initializes the mbarrier pointer for two cta deallocation case. + This initializer sets up the allocator's state, including the shared memory (smem) pointer + holding the base address of the allocated tensor memory, barrier synchronization for + retrieving the tensor memory pointer, allocator warp ID, whether the allocator is being used + for a 2-SM configuration, number of allocated columns in tensor + memory, and the optional mbarrier pointer for deallocation in the 2-SM case. + + If `is_two_cta` is set to True, this will initialize the mbarrier pointer required for tensor + memory deallocation across two CTAs. + + :param alloc_result_dst_smem_ptr: The shared memory pointer that holds the base address of allocated tensor memory. + :type alloc_result_dst_smem_ptr: cute.Pointer + :param barrier_for_retrieve: The named barrier for retrieving the tensor memory pointer. + :type barrier_for_retrieve: pipeline.NamedBarrier + :param allocator_warp_id: The warp ID of the allocator warp, defaults to 0. + :type allocator_warp_id: int, optional + :param is_two_cta: Whether the allocator should coordinate two CTAs, defaults to False. + :type is_two_cta: bool, optional + :param num_allocated_columns: The number of columns allocated in tensor memory, defaults to 0. + :type num_allocated_columns: int, optional + :param two_cta_tmem_dealloc_mbar_ptr: The mbarrier pointer required for two-CTA tensor memory deallocation, optional. + :type two_cta_tmem_dealloc_mbar_ptr: cute.Pointer, optional + :param loc: Optional codegen location for debugging and error reporting. + :type loc: Any, optional + :param ip: Optional insertion point for codegen. + :type ip: Any, optional + + :raises AssertionError: If two_cta_tmem_dealloc_mbar_ptr is None while is_two_cta is True. """ # TODO: automatically maintain a smem address self._alloc_result_dst_smem_ptr = alloc_result_dst_smem_ptr @@ -83,7 +119,7 @@ class TmemAllocator: # Init tmem dealloc mbarrier if two cta if const_expr(self._is_two_cta): - self._init_dealloc_mbarrier() + self._init_dealloc_mbarrier(loc=loc, ip=ip) def __extract_mlir_values__(self) -> list[ir.Value]: values = extract_mlir_values(self._alloc_result_dst_smem_ptr) @@ -137,8 +173,9 @@ class TmemAllocator: return False return True + @dsl_user_op @cute.jit - def allocate(self, num_columns: int): + def allocate(self, num_columns: int, *, loc=None, ip=None): """Allocate a block of tensor memory. This method allocates a block of tensor memory from allocator warp and returns a handle to retrieve @@ -152,29 +189,34 @@ class TmemAllocator: "total allocated columns must be less than or equal to 512" ) - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) + warp_idx = cute.arch.warp_idx(loc=loc, ip=ip) + warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip) _is_allocator_warp = warp_idx == self._allocator_warp_id if _is_allocator_warp: cute.arch.alloc_tmem( num_columns, self._alloc_result_dst_smem_ptr, is_two_cta=self._is_two_cta, + loc=loc, + ip=ip, ) self._num_allocated_columns += num_columns - @cute.jit - def wait_for_alloc(self): + @dsl_user_op + def wait_for_alloc(self, *, loc=None, ip=None): """Wait for the allocator warp to finish allocation. This method is used to synchronize the allocator warp with the other warps before retrieving tmem ptr. """ - self._barrier_for_retrieve.arrive_and_wait() + self._barrier_for_retrieve.arrive_and_wait(loc=loc, ip=ip) - @cute.jit + @dsl_user_op def retrieve_ptr( self, dtype: Type[Numeric] = Float32, + *, + loc=None, + ip=None, ) -> cute.Pointer: """Retrieve the pointer to the allocated tensor memory. @@ -185,30 +227,36 @@ class TmemAllocator: dtype, alignment=16, ptr_to_buffer_holding_addr=self._alloc_result_dst_smem_ptr, + loc=loc, + ip=ip, ) + @dsl_user_op @cute.jit - def relinquish_alloc_permit(self): + def relinquish_alloc_permit(self, *, loc=None, ip=None): """Relinquish the tensor memory allocation permit. This method relinquishes the tensor memory allocation permit for the allocator warp, promising the allocator warp will not allocate any more tensor memory. """ - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) + warp_idx = cute.arch.warp_idx(loc=loc, ip=ip) + warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip) _is_allocator_warp = warp_idx == self._allocator_warp_id if _is_allocator_warp: - cute.arch.relinquish_tmem_alloc_permit(is_two_cta=self._is_two_cta) + cute.arch.relinquish_tmem_alloc_permit( + is_two_cta=self._is_two_cta, loc=loc, ip=ip + ) + @dsl_user_op @cute.jit - def free(self, tmem_ptr: cute.Pointer, num_columns: int = 0): + def free(self, tmem_ptr: cute.Pointer, num_columns: int = 0, *, loc=None, ip=None): """Deallocate the tensor memory. This method sync on mbarrier (for two cta use case) and deallocates the tensor memory from the allocator warp. User can optionally specify the number of columns to deallocate. If not specified, all allocated columns will be deallocated. """ - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) + warp_idx = cute.arch.warp_idx(loc=loc, ip=ip) + warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip) _is_allocator_warp = warp_idx == self._allocator_warp_id assert num_columns <= self._num_allocated_columns, ( @@ -223,15 +271,67 @@ class TmemAllocator: self._num_allocated_columns -= num_deallocate_columns if _is_allocator_warp: if const_expr(self._is_two_cta): + bid_in_cluster = cute.arch.block_idx_in_cluster(loc=loc, ip=ip) _cta_rank_in_cluster = cute.arch.make_warp_uniform( - cute.arch.block_idx_in_cluster() + bid_in_cluster, loc=loc, ip=ip ) # Arrive and wait for dealloc signal from peer cta cute.arch.mbarrier_arrive( - self._two_cta_tmem_dealloc_mbar_ptr, _cta_rank_in_cluster ^ 1 + self._two_cta_tmem_dealloc_mbar_ptr, + _cta_rank_in_cluster ^ 1, + loc=loc, + ip=ip, + ) + cute.arch.mbarrier_wait( + self._two_cta_tmem_dealloc_mbar_ptr, 0, loc=loc, ip=ip ) - cute.arch.mbarrier_wait(self._two_cta_tmem_dealloc_mbar_ptr, 0) # Deallocate tmem cute.arch.dealloc_tmem( - tmem_ptr, num_deallocate_columns, is_two_cta=self._is_two_cta + tmem_ptr, + num_deallocate_columns, + is_two_cta=self._is_two_cta, + loc=loc, + ip=ip, ) + + +# Set explicit signature for Sphinx documentation to avoid issues with @dsl_user_op decorator +TmemAllocator.__init__.__signature__ = inspect.Signature( + [ + inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD), + inspect.Parameter( + "alloc_result_dst_smem_ptr", + inspect.Parameter.POSITIONAL_OR_KEYWORD, + annotation=cute.Pointer, + ), + inspect.Parameter( + "barrier_for_retrieve", + inspect.Parameter.POSITIONAL_OR_KEYWORD, + annotation=pipeline.NamedBarrier, + ), + inspect.Parameter( + "allocator_warp_id", + inspect.Parameter.POSITIONAL_OR_KEYWORD, + default=0, + annotation=int, + ), + inspect.Parameter( + "is_two_cta", + inspect.Parameter.POSITIONAL_OR_KEYWORD, + default=False, + annotation=bool, + ), + inspect.Parameter( + "num_allocated_columns", + inspect.Parameter.POSITIONAL_OR_KEYWORD, + default=0, + annotation=int, + ), + inspect.Parameter( + "two_cta_tmem_dealloc_mbar_ptr", + inspect.Parameter.POSITIONAL_OR_KEYWORD, + default=None, + annotation=Optional[cute.Pointer], + ), + ] +) diff --git a/python/CuTeDSL/requirements.txt b/python/CuTeDSL/requirements.txt index 75eb76cb..19e72110 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.3.0.dev0 +nvidia-cutlass-dsl==4.3.0 diff --git a/python/cutlass_cppgen/__init__.py b/python/cutlass_cppgen/__init__.py index 2491da63..9bdd259c 100644 --- a/python/cutlass_cppgen/__init__.py +++ b/python/cutlass_cppgen/__init__.py @@ -133,7 +133,7 @@ def get_option_registry(): this._option_registry = OptionRegistry(device_cc()) return this._option_registry -this.__version__ = '4.3.0' +this.__version__ = '4.2.1' from cutlass_cppgen.backend import create_memory_pool from cutlass_cppgen.emit.pytorch import pytorch diff --git a/python/cutlass_cppgen/backend/evt/backend/sm90_emitter.py b/python/cutlass_cppgen/backend/evt/backend/sm90_emitter.py index 3c058aa8..3dcf4eff 100644 --- a/python/cutlass_cppgen/backend/evt/backend/sm90_emitter.py +++ b/python/cutlass_cppgen/backend/evt/backend/sm90_emitter.py @@ -90,7 +90,7 @@ class Sm90Emitter: tile_description=operation.tile_description, schedule=operation.tile_description.epilogue_schedule, element_c=operation.C.element, - element_d=operation.C.element, + element_d=operation.D.element, fusion_callbacks=fusion_callbacks ) diff --git a/python/cutlass_cppgen/backend/evt/frontend/python_ast.py b/python/cutlass_cppgen/backend/evt/frontend/python_ast.py index 8727b754..1d45d7a3 100644 --- a/python/cutlass_cppgen/backend/evt/frontend/python_ast.py +++ b/python/cutlass_cppgen/backend/evt/frontend/python_ast.py @@ -140,7 +140,7 @@ class PythonASTFrontend(EVTFrontendBase, ast.NodeVisitor): self.add_edge(rhs, name, weight=1) return name - def visit_Assign(self, node: ast.BinOp): + def visit_Assign(self, node: ast.Assign): target = self.visit(node.targets[0]) value = self.visit(node.value) # Create the assign node diff --git a/python/cutlass_cppgen/backend/evt/ir/layout_algorithm.py b/python/cutlass_cppgen/backend/evt/ir/layout_algorithm.py index 9d453b1f..ddfbec88 100644 --- a/python/cutlass_cppgen/backend/evt/ir/layout_algorithm.py +++ b/python/cutlass_cppgen/backend/evt/ir/layout_algorithm.py @@ -119,8 +119,8 @@ def _get_first_rhs_nonzero_stride(stride_list, idx): for i in range(idx+1, len(stride_list)): if stride_list[i] != 0: return i - else: - return None + else: + return None def reshape(layout, new_shape): """ diff --git a/python/cutlass_cppgen/backend/evt/passes/graph_drawer.py b/python/cutlass_cppgen/backend/evt/passes/graph_drawer.py index 8a28c6e4..43e40343 100644 --- a/python/cutlass_cppgen/backend/evt/passes/graph_drawer.py +++ b/python/cutlass_cppgen/backend/evt/passes/graph_drawer.py @@ -99,11 +99,10 @@ class EVTGraphDrawer: stride = node.tensor.stride label += f"|shape={shape}|stride={stride}" - if hasattr(node, "store_tensor"): - if node.store_tensor is not None: - store_shape = node.store_tensor.shape - store_stride = node.store_tensor.stride - label += f"|store_shape={store_shape}|stride_stride={store_stride}" + if hasattr(node, "store_tensor") and node.store_tensor is not None: + store_shape = node.store_tensor.shape + store_stride = node.store_tensor.stride + label += f"|store_shape={store_shape}|store_stride={store_stride}" label += "}" return label @@ -114,7 +113,7 @@ class EVTGraphDrawer: name: str ): import pydot - dot_graph = pydot.Dot(name, randir="TB") + dot_graph = pydot.Dot(name, rankdir="TB") for node in graph.nodes_meta: style = self._get_node_style(node) label = self._get_node_label(node) @@ -133,11 +132,11 @@ class EVTGraphDrawer: return dot_graph - def get_dot_graph(self) -> pydot.Dot: + def get_dot_graph(self) -> "pydot.Dot": return [(key, self.get_dot_graph_by_name(key)) for key in self._dot_graphs.keys()] - def get_dot_graph_by_name(self, name) -> pydot.Dot: + def get_dot_graph_by_name(self, name) -> "pydot.Dot": return self._dot_graphs[name] - def get_main_dot_graph(self) -> pydot.Dot: + def get_main_dot_graph(self) -> "pydot.Dot": return self._dot_graphs[self._name] diff --git a/python/cutlass_cppgen/backend/evt/passes/pass_preprocess_red.py b/python/cutlass_cppgen/backend/evt/passes/pass_preprocess_red.py index 6423a2b8..3db839ac 100644 --- a/python/cutlass_cppgen/backend/evt/passes/pass_preprocess_red.py +++ b/python/cutlass_cppgen/backend/evt/passes/pass_preprocess_red.py @@ -51,15 +51,14 @@ class PassPreprocessRed(EVTPassBase): # Step 1: find the compute nodes with op=red red_compute_nodes = [] for node_meta in self.dag_ir.nodes_meta: - if isinstance(node_meta, ComputeNode): - if type(node_meta.fn) == tuple: - # To keep the frontend simple, the reduction nodes - # are parsed into compute nodes by default - # The simple heuristic to distinguish between compute - # and reduction node is that compute node is a single function, - # while the reduction node is a tuple of functions for - # in-register reduction and atomic global memory reduction - red_compute_nodes.append(node_meta.name) + if isinstance(node_meta, ComputeNode) and type(node_meta.fn) == tuple: + # To keep the frontend simple, the reduction nodes + # are parsed into compute nodes by default + # The simple heuristic to distinguish between compute + # and reduction node is that compute node is a single function, + # while the reduction node is a tuple of functions for + # in-register reduction and atomic global memory reduction + red_compute_nodes.append(node_meta.name) # Step 2: for each compute, merge it with the succeeding store for node in red_compute_nodes: diff --git a/python/cutlass_library/gemm_operation.py b/python/cutlass_library/gemm_operation.py index b9c8751d..7bbdcdb2 100644 --- a/python/cutlass_library/gemm_operation.py +++ b/python/cutlass_library/gemm_operation.py @@ -77,6 +77,8 @@ class GemmOperation: GemmKind.BlockwiseUniversal3x, GemmKind.GroupedBlockwiseUniversal3x, GemmKind.BlockScaledSparseUniversal3x, + GemmKind.MoeGroupedUniversal3x, + GemmKind.BlockScaledMoeGroupedUniversal3x, } self.is_3x = gemm_kind in kinds_3x self.prefix = "3x" if self.is_3x else "" @@ -927,8 +929,14 @@ ${compile_guard_end} gemm_shape_type = "cute::Shape" grouped_gemm_shape_type = "cute::Shape" grouped_gemm_shape_type = "cutlass::gemm::GroupProblemShape<" + grouped_gemm_shape_type + ">" - - return gemm_shape_type if not is_grouped(operation.gemm_kind) else grouped_gemm_shape_type + moe_gemm_shape_type = "cute::Shape" + moe_gemm_shape_type = "cutlass::gemm::MoEProblemShape<" + moe_gemm_shape_type + ">" + if is_moe(operation.gemm_kind): + return moe_gemm_shape_type + elif is_grouped(operation.gemm_kind): + return grouped_gemm_shape_type + else: + return gemm_shape_type def emit(self, operation): _LOGGER.debug("*** EmitGemmConfigurationLibrary::emit(operation)") @@ -943,6 +951,7 @@ ${compile_guard_end} instruction_shape = operation.tile_description.math_instruction.instruction_shape cluster_m = operation.tile_description.cluster_shape[0] cluster_n = operation.tile_description.cluster_shape[1] + cta_m = tile_shape[0] // cluster_m if cluster_m > 0 else tile_shape[0] cta_n = tile_shape[1] // cluster_n if cluster_n > 0 else tile_shape[1] tile_shape_m, tile_shape_n, tile_shape_k = operation.get_collective_tile_shape() @@ -1023,6 +1032,13 @@ ${compile_guard_end} element_a = f'cute::tuple<{str(element_a)},{str(DataTypeTag[operation.ScaleFactorA])}>' element_b = f'cute::tuple<{str(element_b)},{str(DataTypeTag[operation.ScaleFactorB])}>' + if is_moe(operation.gemm_kind): + if DataTypeSize[operation.A.element] == 4 and operation.ScaleFactorA == DataType.ue4m3: + element_a = f"cutlass::nv_float4_t<{DataTypeTag[operation.A.element]}>" + + if DataTypeSize[operation.B.element] == 4 and operation.ScaleFactorB == DataType.ue4m3: + element_b = f"cutlass::nv_float4_t<{DataTypeTag[operation.B.element] }>" + alignment_c = get_tma_alignment(operation.C.element) \ if is_tma_epilogue(operation.epilogue_schedule) and opcode_class_epi != OpcodeClass.Simt \ else operation.C.alignment @@ -1480,6 +1496,8 @@ class EmitGemmConfigurationLibrary: GemmKind.BlockwiseUniversal3x: EmitGemmUniversal3xInstance, GemmKind.GroupedBlockwiseUniversal3x: EmitGemmUniversal3xInstance, GemmKind.BlockScaledSparseUniversal3x: EmitGemmUniversal3xInstance, + GemmKind.MoeGroupedUniversal3x: EmitGemmUniversal3xInstance, + GemmKind.BlockScaledMoeGroupedUniversal3x: EmitGemmUniversal3xInstance, } self.gemm_kind_wrappers = { @@ -1497,6 +1515,8 @@ class EmitGemmConfigurationLibrary: GemmKind.BlockwiseUniversal3x: 'BlockwiseGemmUniversal3xOperation', GemmKind.GroupedBlockwiseUniversal3x: 'GroupedBlockwiseGemmUniversal3xOperation', GemmKind.BlockScaledSparseUniversal3x: 'BlockScaledSparseGemmUniversal3xOperation', + GemmKind.MoeGroupedUniversal3x: 'MoeGroupedGemmUniversal3xOperation', + GemmKind.BlockScaledMoeGroupedUniversal3x: 'BlockScaledMoeGroupedGemmUniversal3xOperation', } self.wmma_guard_start = "#if defined(CUTLASS_ARCH_WMMA_SM${sm_number}_ENABLED)" diff --git a/python/cutlass_library/generator.py b/python/cutlass_library/generator.py index 34fc7336..d94143fa 100644 --- a/python/cutlass_library/generator.py +++ b/python/cutlass_library/generator.py @@ -5841,7 +5841,24 @@ def GenerateSM90_TensorOp_fp8_WGMMA_gemm_with_blockwise(manifest, cuda_version, level=instantiation_level) tile_descriptions = list() - + tile_descriptions.append( + TileDescription( + threadblock_shape=[ + 256, + 128, + 128 + ], + stages=0, + warp_count=[4, 1, 1], + math_instruction=MathInstruction( + [128, 128, 32], + DataType.e5m2, DataType.e4m3, DataType.f32, + OpcodeClass.TensorOp, + MathOperation.multiply_add), + min_compute=90, + max_compute=90, + cluster_shape=[1,2,1], + explicit_vector_sizes=[1, 128, 128])) for desc in tile_descriptions_: desc.explicit_vector_sizes = [1, desc.tile_shape[1], desc.tile_shape[2]] tile_descriptions.append(copy.deepcopy(desc)) @@ -7043,6 +7060,104 @@ def GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types_mixed, [[kernel_schedule, epi_schedule]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) +def GenerateSM100_TensorOp_16b_UMMA_alignx_gemm(manifest, cuda_version, gemm_kind=GemmKind.Universal3x): + if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): + return + + instantiation_level = manifest.get_instantiation_level(pruned_level=490, default_level=490, exhaustive_level=9999) + + # layouts for ABC and their alignments. C alignment will be set later based on output type + layouts = [ + [[LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 1]], + [[LayoutType.RowMajor, 4], [LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 1]], + [[LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 1]], + [[LayoutType.ColumnMajor, 4], [LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 1]], + [[LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 1]], + [[LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 1]], + [[LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 1]], + [[LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 1]], + ] + + thor_sm = ThorSMRenumbering(cuda_version) + + math_instructions_1sm, _ = generate_16b_math_instructions_sm100(instantiation_level) + + min_cc = 100 + max_cc = 100 + max_cc = max(max_cc, thor_sm) + + grouped = is_grouped(gemm_kind) + if grouped: + return + cluster_shapes_1sm= [[1,1,1]] + + tile_schedulers = [ + TileSchedulerType.Default + ] + + # 1xSM MMA kernels + for math_inst in math_instructions_1sm: + tile_descriptions = [] + for cluster_shape in cluster_shapes_1sm: + multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape + tile_descriptions.append( + TileDescription([ + math_inst.instruction_shape[0] * multiplier_1sm[0], + math_inst.instruction_shape[1] * multiplier_1sm[1], + math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], + 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) + + data_types = [ + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : math_inst.element_accumulator, + "d_type" : math_inst.element_accumulator, + "acc_type" : math_inst.element_accumulator, + "epi_type" : math_inst.element_accumulator, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.void, + "d_type" : math_inst.element_accumulator, + "acc_type" : math_inst.element_accumulator, + "epi_type" : math_inst.element_accumulator, + }, + ] + + kernel_schedule = KernelScheduleType.WarpSpecialized1SmSm100 + epi_schedule = EpilogueScheduleType.NoSmemWarpSpecialized1Sm + CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types, + [[kernel_schedule, epi_schedule]], + tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) + + # for mixed precision kernels, also generate kernels that write output matrix in the A/B format + # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) + if math_inst.element_a != math_inst.element_accumulator: + data_types_mixed = [ + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : math_inst.element_a, + "d_type" : math_inst.element_a, + "acc_type" : math_inst.element_accumulator, + "epi_type" : math_inst.element_accumulator, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.void, + "d_type" : math_inst.element_a, + "acc_type" : math_inst.element_accumulator, + "epi_type" : math_inst.element_accumulator, + }, + ] + + CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types_mixed, + [[kernel_schedule, epi_schedule]], + tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) + def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.Universal3x): if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return @@ -7198,16 +7313,19 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] - - for data_type in data_types: - if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\ - ( data_type["d_type"] == DataType.e5m2 ): - continue - kernel_schedule = to_grouped_schedule(KernelScheduleType.TmaWarpSpecialized1SmSm100, grouped) - epi_schedule = to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized1Sm, grouped) - CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, - [[kernel_schedule, epi_schedule]], - tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) + for tile_description in tile_descriptions: + for layout in layouts: + for data_type in data_types: + if layout[1][0] == LayoutType.RowMajor and tile_description.math_instruction.instruction_shape[1] % 16 != 0: + continue + if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\ + ( data_type["d_type"] == DataType.e5m2 ): + continue + kernel_schedule = to_grouped_schedule(KernelScheduleType.TmaWarpSpecialized1SmSm100, grouped) + epi_schedule = to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized1Sm, grouped) + CreateGemmUniversal3xOperator(manifest, [layout], [tile_description], data_type, + [[kernel_schedule, epi_schedule]], + tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) # 2xSM MMA kernels @@ -7341,6 +7459,158 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, [[kernel_schedule, epi_schedule]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) +def GenerateSM100_TensorOp_fp8_UMMA_alignx_gemm(manifest, cuda_version, gemm_kind=GemmKind.Universal3x): + if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): + return + + instantiation_level = manifest.get_instantiation_level(pruned_level=591 , default_level=591 , exhaustive_level=9999) + + # layouts for ABC and their alignments. + layouts = [ + [[LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 1]], # TN Layout + [[LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 1]], # TN Layout + ] + + thor_sm = ThorSMRenumbering(cuda_version) + + min_cc = 100 + max_cc = 100 + max_cc = max(max_cc, thor_sm) + + epi_type = DataType.f32 + grouped = is_grouped(gemm_kind) + + math_instructions_1sm, _ = generate_fp8_math_instructions_sm100(instantiation_level, enable_runtime_dtype=not grouped) + + cluster_shapes_1sm = [[1,1,1]] + + tile_schedulers = [ + TileSchedulerType.Default + ] + + # 1xSM MMA kernels + for math_inst in math_instructions_1sm: + tile_descriptions = [] + for cluster_shape in cluster_shapes_1sm: + multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape + tile_descriptions.append( + TileDescription([ + math_inst.instruction_shape[0] * multiplier_1sm[0], + math_inst.instruction_shape[1] * multiplier_1sm[1], + math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], + 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) + + data_types = [ + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.f16, + "d_type" : DataType.f16, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.f16, + "d_type" : DataType.e4m3, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.f16, + "d_type" : DataType.e5m2, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.bf16, + "d_type" : DataType.bf16, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.bf16, + "d_type" : DataType.e4m3, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.bf16, + "d_type" : DataType.e5m2, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.f32, + "d_type" : DataType.f32, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.void, + "d_type" : DataType.f16, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.void, + "d_type" : DataType.bf16, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.void, + "d_type" : DataType.f32, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.void, + "d_type" : DataType.e4m3, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.void, + "d_type" : DataType.e5m2, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + } + ] + + + for data_type in data_types: + if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\ + ( data_type["d_type"] == DataType.e5m2 ): + continue + + kernel_schedule = KernelScheduleType.WarpSpecialized1SmSm100 + epi_schedule = EpilogueScheduleType.NoSmemWarpSpecialized1Sm + CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, + [[kernel_schedule, epi_schedule]], + tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) + def GenerateSM100_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version, gemm_kind=GemmKind.BlockwiseUniversal3x): if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return @@ -7467,7 +7737,7 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version, continue is_runtime_datatype_a = is_runtime_datatype(data_type["a_type"]) - is_runtime_datatype_b = is_runtime_datatype(data_type["d_type"]) + is_runtime_datatype_b = is_runtime_datatype(data_type["b_type"]) # A/B datatypes should be both static or dynamic if (is_runtime_datatype_a != is_runtime_datatype_b): @@ -8585,6 +8855,418 @@ def GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifes CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, schedules , tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind ) +def GenerateSM100_TensorOp_16b_UMMA_moe_gemm(manifest, cuda_version, gemm_kind=GemmKind.MoeGroupedUniversal3x): + # SM100 MOE GEMM + if not CudaToolkitVersionSatisfies(cuda_version, 13, 0): + return + + instantiation_level = manifest.get_instantiation_level(pruned_level=490, default_level=490, exhaustive_level=9999) + # layouts for ABC and their alignments. C alignment will be set later based on output type + layouts = [ + [[LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 0]], + ] + + thor_sm = ThorSMRenumbering(cuda_version) + + math_instructions_1sm, math_instructions_2sm = generate_16b_math_instructions_sm100(instantiation_level) + + min_cc = 100 + max_cc = 100 + max_cc = max(max_cc, thor_sm) + + cluster_shapes= [[1,1,1]] + + tile_schedulers = [ + TileSchedulerType.Default + ] + # 1xSM MMA kernels + for math_inst in math_instructions_1sm: + tile_descriptions = [] + for cluster_shape in cluster_shapes: + multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape + tile_descriptions.append( + TileDescription([ + math_inst.instruction_shape[0] * multiplier_1sm[0], + math_inst.instruction_shape[1] * multiplier_1sm[1], + math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], + 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) + data_types = [ + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : math_inst.element_accumulator, + "d_type" : math_inst.element_b, + "acc_type" : math_inst.element_accumulator, + "epi_type" : math_inst.element_accumulator, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.void, + "d_type" : math_inst.element_b, + "acc_type" : math_inst.element_accumulator, + "epi_type" : math_inst.element_accumulator, + }, + ] + # Set alignment d based on Destination format. + for layout in layouts: + layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] + kernel_schedule = KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmSm100 + epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1Sm + CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types, + [[kernel_schedule, epi_schedule]], + tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) + +def GenerateSM100_TensorOp_fp8_UMMA_moe_gemm(manifest, cuda_version, gemm_kind=GemmKind.MoeGroupedUniversal3x): + # SM100 MOE GEMM + if not CudaToolkitVersionSatisfies(cuda_version, 13, 0): + return + + instantiation_level = manifest.get_instantiation_level(pruned_level=490, default_level=490, exhaustive_level=9999) + # layouts for ABC and their alignments. C alignment will be set later based on output type + layouts = [ + [[LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 0]], + ] + + thor_sm = ThorSMRenumbering(cuda_version) + + math_instructions_1sm, math_instructions_2sm = generate_fp8_math_instructions_sm100(instantiation_level) + + min_cc = 100 + max_cc = 100 + max_cc = max(max_cc, thor_sm) + # only support 1x1x1 cluster shape + cluster_shapes= [[1,1,1]] + epi_type = DataType.f32 + tile_schedulers = [ + TileSchedulerType.Default + ] + # 1xSM MMA kernels + for math_inst in math_instructions_1sm: + tile_descriptions = [] + for cluster_shape in cluster_shapes: + multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape + tile_descriptions.append( + TileDescription([ + math_inst.instruction_shape[0] * multiplier_1sm[0], + math_inst.instruction_shape[1] * multiplier_1sm[1], + math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], + 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) + data_types = [ + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.f16, + "d_type" : DataType.f16, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.f16, + "d_type" : DataType.e4m3, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.f16, + "d_type" : DataType.e5m2, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.bf16, + "d_type" : DataType.bf16, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.bf16, + "d_type" : DataType.e4m3, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + }, + ] + # Set alignment d based on Destination format. + for layout in layouts: + layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] + kernel_schedule = KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmSm100 + epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1Sm + for data_type in data_types: + if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\ + ( data_type["d_type"] == DataType.e5m2 ): + continue + CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, + [[kernel_schedule, epi_schedule]], + tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) + +def GenerateSM100_TensorOp_mixed_8bits_UMMA_moe_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledMoeGroupedUniversal3x): + # SM100 moe GEMM with mixed F4/F6/F8 inputs + block scale + if not CudaToolkitVersionSatisfies(cuda_version, 13, 0): + return + + instantiation_level = manifest.get_instantiation_level(pruned_level=590, default_level=590, exhaustive_level=9999) + + grouped = is_grouped(gemm_kind) + + layouts = [ + [[LayoutType.RowMajor, 0], [LayoutType.ColumnMajor, 0], [LayoutType.ColumnMajor, 0]], + ] + + math_instructions_1sm, math_instructions_2sm = generate_mxf8f6f4_math_instructions_sm100(instantiation_level, enable_runtime_dtype=not grouped) + + + cluster_shapes_1sm = [[1,1,1]] + + acc_types = [ DataType.f32 ] + + def tile_schedulers(sfdtype): + # Only use the stream-K scheduler for non-void SFD to limit kernel count. When SFD is void, + # the epilogue is the traditional linear combination, for which we already have tests with stream-K. + if sfdtype["type"] == DataType.void or grouped: + return [TileSchedulerType.Default] + else: + return [TileSchedulerType.Default, TileSchedulerType.StreamK] + + thor_sm = ThorSMRenumbering(cuda_version) + + min_cc = 100 + max_cc = 100 + max_cc = max(max_cc, thor_sm) + + epi_type = DataType.f32 + + is_runtime_datatype = lambda runtime_datatype: runtime_datatype in (DataType.f4, DataType.f6, DataType.f8) + + + # 1xSM MMA kernels + for math_inst in math_instructions_1sm: + assert math_inst.opcode_class == OpcodeClass.BlockScaledTensorOp + if DataTypeSize[math_inst.element_a] != DataTypeSize[math_inst.element_b] or DataTypeSize[math_inst.element_a] == 6: + continue + tile_descriptions = [] + for cluster_shape in cluster_shapes_1sm: + multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape + tile_descriptions.append( + TileDescription([ + math_inst.instruction_shape[0] * multiplier_1sm[0], + math_inst.instruction_shape[1] * multiplier_1sm[1], + math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], + 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) + + data_types = [ + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.void, + "d_type" : DataType.bf16, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + "sf_type" : math_inst.element_scale_factor, + "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.void, + "d_type" : DataType.e5m2, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + "sf_type" : math_inst.element_scale_factor, + "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.f16, + "d_type" : DataType.e5m2, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + "sf_type" : math_inst.element_scale_factor, + "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.f16, + "d_type" : DataType.e3m2, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + "sf_type" : math_inst.element_scale_factor, + "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} + }] + + + kernel_schedule = KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100 + epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1Sm + for data_type in data_types: + # Set alignment d based on Destination format. + for layout in layouts: + layout[0][1] = 128 // DataTypeSize[data_type["a_type"]] + layout[1][1] = 128 // DataTypeSize[data_type["b_type"]] + layout[2][1] = 128 // DataTypeSize[data_type["d_type"]] + CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, + schedules=[[kernel_schedule, epi_schedule]], tile_schedulers=[TileSchedulerType.Default], gemm_kind=gemm_kind) + +def GenerateSM100_TensorOp_fp4_UMMA_MoE_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledMoeGroupedUniversal3x): + # SM100 MoE GEMM with F4 + block scale + if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): + return + + instantiation_level = manifest.get_instantiation_level(pruned_level=591, default_level=591, exhaustive_level=9999) + + + # layouts for ABC and their alignments. + layouts = [ + [[LayoutType.RowMajor, 32], [LayoutType.ColumnMajor, 32], [LayoutType.ColumnMajor, 0]], + ] + + shapes_1sm = [ + (128, 64, 64), (128, 128, 64), (128, 192, 64), (128, 256, 64) + ] + math_instructions_1sm = [] + for shape in shapes_1sm: + math_instructions_1sm.append( + MathInstruction( + shape, + DataType.e2m1, DataType.e2m1, DataType.f32, + OpcodeClass.BlockScaledTensorOp, + MathOperation.multiply_add, + DataType.ue8m0) + ) + cluster_shapes_1sm = [[1,1,1]] + + acc_types = [ DataType.f32 ] # Accumulator is always 32 bits for block scaled MMA instructions + + + thor_sm = ThorSMRenumbering(cuda_version) + + min_cc = 100 + max_cc = 100 + max_cc = max(max_cc, thor_sm) + + epi_type = DataType.f32 + + is_runtime_datatype = lambda runtime_datatype: runtime_datatype in (DataType.f4, DataType.f6, DataType.f8) + + + # 1xSM MMA kernels + for math_inst in math_instructions_1sm: + assert math_inst.opcode_class == OpcodeClass.BlockScaledTensorOp + + tile_descriptions = [] + for cluster_shape in cluster_shapes_1sm: + multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape + tile_descriptions.append( + TileDescription([ + math_inst.instruction_shape[0] * multiplier_1sm[0], + math_inst.instruction_shape[1] * multiplier_1sm[1], + math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], + 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) + assert math_inst.instruction_shape[2] * 4 == 256 + + data_types = [ + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.bf16, + "d_type" : DataType.bf16, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + "sf_type" : math_inst.element_scale_factor, + "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.void, + "d_type" : DataType.e2m1, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + "sf_type" : math_inst.element_scale_factor, + "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.void, + "d_type" : DataType.e5m2, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + "sf_type" : math_inst.element_scale_factor, + "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.f16, + "d_type" : DataType.e5m2, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + "sf_type" : math_inst.element_scale_factor, + "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.void, + "d_type" : DataType.e2m1, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + "sf_type" : math_inst.element_scale_factor, + "sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor} + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.f16, + "d_type" : DataType.e2m1, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + "sf_type" : math_inst.element_scale_factor, + "sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor} + }, + { + "a_type" : math_inst.element_a, + "b_type" : math_inst.element_b, + "c_type" : DataType.f16, + "d_type" : DataType.e2m1, + "acc_type" : math_inst.element_accumulator, + "epi_type" : epi_type, + "sf_type" : math_inst.element_scale_factor, + "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} + } + ] + + # Set alignment d based on Destination format. + for layout in layouts: + layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] + + for layout in layouts: + for data_type in data_types: + if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.RowMajor): + data_type["sfd_type"]["layout"] = layout[2][0] # For FP4 output , the scalefactor layout is same layout as D layout. + if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.ColumnMajor): + continue + + # E2M1 x E2M1, vector size 32, E8 + # E2M1 x E2M1, vector size 16, UE4M3 + isFp4 = math_inst.element_scale_factor == DataType.ue8m0 and math_inst.element_a == DataType.e2m1 and math_inst.element_b == DataType.e2m1 + + kernel_schedule = KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100 + epi_schedule = EpilogueScheduleType.ScheduleAuto + + CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, [[kernel_schedule, epi_schedule]] + , tile_schedulers=[TileSchedulerType.Default], gemm_kind=gemm_kind + ) + def GenerateSM103_TensorOp_fp4_ultra_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledUniversal3x): # SM100 MMA with F4 + block scale @@ -11047,17 +11729,22 @@ def GenerateSM100(manifest, cuda_version): # Dense Gemm # GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version) - + GenerateSM100_TensorOp_16b_UMMA_alignx_gemm(manifest, cuda_version) GenerateSM100_TensorOp_32b_UMMA_gemm(manifest, cuda_version) if not bool(set(manifest.compute_capabilities_feature_set).intersection(arch_family_cc)): GenerateSM100_TensorOp_int8_UMMA_gemm(manifest, cuda_version) GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version) + GenerateSM100_TensorOp_fp8_UMMA_alignx_gemm(manifest, cuda_version) # grouped GEMM GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.GroupedUniversal3x) GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.GroupedUniversal3x) - + # MOE grouped GEMM + GenerateSM100_TensorOp_16b_UMMA_moe_gemm(manifest, cuda_version) + GenerateSM100_TensorOp_fp8_UMMA_moe_gemm(manifest, cuda_version) + GenerateSM100_TensorOp_mixed_8bits_UMMA_moe_gemm_with_block_scaled(manifest, cuda_version) + GenerateSM100_TensorOp_fp4_UMMA_MoE_gemm_with_block_scaled(manifest, cuda_version) # StreamK is included in regular generation GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version) diff --git a/python/cutlass_library/library.py b/python/cutlass_library/library.py index 3440360d..759c6583 100644 --- a/python/cutlass_library/library.py +++ b/python/cutlass_library/library.py @@ -322,7 +322,8 @@ def is_complex(data_type): return False def is_block_scaled(gemm_kind): - return gemm_kind in (GemmKind.BlockScaledUniversal3x, GemmKind.GroupedBlockScaledUniversal3x, GemmKind.BlockScaledSparseUniversal3x) + return gemm_kind in (GemmKind.BlockScaledUniversal3x, GemmKind.GroupedBlockScaledUniversal3x, GemmKind.BlockScaledSparseUniversal3x, + GemmKind.BlockScaledMoeGroupedUniversal3x) def is_blockwise(gemm_kind): return gemm_kind in (GemmKind.BlockwiseUniversal3x, GemmKind.GroupedBlockwiseUniversal3x) @@ -331,6 +332,8 @@ def is_grouped(gemm_kind): return gemm_kind in (GemmKind.GroupedUniversal3x, GemmKind.GroupedBlockScaledUniversal3x, GemmKind.GroupedBlockwiseUniversal3x) +def is_moe(gemm_kind): + return gemm_kind in (GemmKind.MoeGroupedUniversal3x, GemmKind.BlockScaledMoeGroupedUniversal3x) # def get_complex_from_real(real_type): for r, c in RealComplexBijection: @@ -513,6 +516,8 @@ class KernelScheduleType(enum.Enum): TmaWarpSpecialized1SmSm100 = enum_auto() TmaWarpSpecialized2SmSm100 = enum_auto() + WarpSpecialized1SmSm100 = enum_auto() + ImplicitTmaWarpSpecialized1SmSm100 = enum_auto() ImplicitTmaWarpSpecialized2SmSm100 = enum_auto() @@ -528,6 +533,9 @@ class KernelScheduleType(enum.Enum): PtrArrayMxf8f6f4TmaWarpSpecialized1SmSm100 = enum_auto() PtrArrayMxf8f6f4TmaWarpSpecialized2SmSm100 = enum_auto() + MixedTmaCpAsyncWarpSpecialized1SmSm100 = enum_auto() + MixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100 = enum_auto() + SparseTmaWarpSpecialized1SmSm100 = enum_auto() SparseTmaWarpSpecialized2SmSm100 = enum_auto() @@ -621,6 +629,7 @@ KernelScheduleTag = { KernelScheduleType.TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmSm100', KernelScheduleType.TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmSm100', + KernelScheduleType.WarpSpecialized1SmSm100: 'cutlass::gemm::KernelWarpSpecialized1SmSm100', KernelScheduleType.ImplicitTmaWarpSpecialized1SmSm100: 'cutlass::conv::KernelImplicitTmaWarpSpecialized1SmSm100', KernelScheduleType.ImplicitTmaWarpSpecialized2SmSm100: 'cutlass::conv::KernelImplicitTmaWarpSpecialized2SmSm100', @@ -641,7 +650,8 @@ KernelScheduleTag = { KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedBlockwise1SmSm100', KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedBlockwise2SmSm100', - + KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmSm100: 'cutlass::gemm::KernelMixedTmaCpAsyncWarpSpecialized1SmSm100', + KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100: 'cutlass::gemm::KernelMixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100', KernelScheduleType.Mxf4TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmMxf4Sm100', KernelScheduleType.Mxf4TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmMxf4Sm100', KernelScheduleType.Nvf4TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmNvf4Sm100', @@ -738,6 +748,7 @@ KernelScheduleSuffixes = { KernelScheduleType.TmaWarpSpecialized1SmSm100: '_1sm', KernelScheduleType.TmaWarpSpecialized2SmSm100: '_2sm', + KernelScheduleType.WarpSpecialized1SmSm100: '_cpasync_1sm', KernelScheduleType.ImplicitTmaWarpSpecialized1SmSm100: '_1sm', KernelScheduleType.ImplicitTmaWarpSpecialized2SmSm100: '_2sm', @@ -758,6 +769,9 @@ KernelScheduleSuffixes = { KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecialized1SmSm100: '_1sm', KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecialized2SmSm100: '_2sm', + KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmSm100: '_mixed_cpasync_1sm', + KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100: '_mixed_cpasync_1sm', + KernelScheduleType.Mxf4TmaWarpSpecialized1SmSm100: '_o_vs32_1sm', KernelScheduleType.Mxf4TmaWarpSpecialized2SmSm100: '_o_vs32_2sm', KernelScheduleType.Nvf4TmaWarpSpecialized1SmSm100: '_o_vs16_1sm', @@ -1208,6 +1222,9 @@ class GemmKind(enum.Enum): BlockwiseUniversal3x = enum_auto() GroupedBlockwiseUniversal3x = enum_auto() BlockScaledSparseUniversal3x = enum_auto() + MoeGroupedUniversal3x = enum_auto() + BlockScaledMoeGroupedUniversal3x = enum_auto() + # GemmKindNames = { @@ -1224,7 +1241,9 @@ GemmKindNames = { GemmKind.GroupedBlockScaledUniversal3x: "gemm_grouped", GemmKind.BlockwiseUniversal3x: "gemm", GemmKind.GroupedBlockwiseUniversal3x: "gemm_grouped", - GemmKind.BlockScaledSparseUniversal3x: "spgemm" + GemmKind.BlockScaledSparseUniversal3x: "spgemm", + GemmKind.MoeGroupedUniversal3x: "moe_gemm", + GemmKind.BlockScaledMoeGroupedUniversal3x: "moe_gemm", } # diff --git a/python/setup_cutlass.py b/python/setup_cutlass.py index 8b53d8f4..acc0c46e 100644 --- a/python/setup_cutlass.py +++ b/python/setup_cutlass.py @@ -51,7 +51,7 @@ setup_pycute.perform_setup() setup( name='cutlass_cppgen', - version='4.3.0', + version='4.2.0', description='CUTLASS Pythonic Interface', package_dir={'': '.'}, packages=[ diff --git a/python/setup_library.py b/python/setup_library.py index 4ee11af9..c56d6b55 100644 --- a/python/setup_library.py +++ b/python/setup_library.py @@ -36,7 +36,7 @@ from setuptools import setup def perform_setup(): setup( name='cutlass_library', - version='4.3.0', + version='4.2.1', description='CUTLASS library generation scripts', packages=['cutlass_library'] ) diff --git a/python/setup_pycute.py b/python/setup_pycute.py index 37ba01cb..0bad050f 100644 --- a/python/setup_pycute.py +++ b/python/setup_pycute.py @@ -36,7 +36,7 @@ from setuptools import setup def perform_setup(): setup( name='pycute', - version='4.3.0', + version='4.2.1', description='Python implementation of CuTe', packages=['pycute'], ) diff --git a/test/unit/nvrtc/CMakeLists.txt b/test/unit/nvrtc/CMakeLists.txt index 70c50a4a..85a1471c 100644 --- a/test/unit/nvrtc/CMakeLists.txt +++ b/test/unit/nvrtc/CMakeLists.txt @@ -41,6 +41,7 @@ macro(add_nvrtc_headers BASE_DIR FILES) -DFILE_IN="${BASE_DIR}/${CUTLASS_FILE}" -DFILE_OUT="${OUTPUT_FILE}" -DVARIABLE_NAME="${VARIABLE_NAME}" + -DMSVC=${MSVC} -P ${PROJECT_SOURCE_DIR}/bin2hex.cmake DEPENDS ${BASE_DIR}/${CUTLASS_FILE} ) diff --git a/test/unit/transform/device/CMakeLists.txt b/test/unit/transform/device/CMakeLists.txt index d5c0d5ee..3999b8a3 100644 --- a/test/unit/transform/device/CMakeLists.txt +++ b/test/unit/transform/device/CMakeLists.txt @@ -56,3 +56,50 @@ cutlass_test_unit_add_executable( sm90_sparse_gemm_compressor_f8.cu ) +add_custom_target( + cutlass_test_unit_sm100_structured_sparse_gemm_compressor + DEPENDS + cutlass_test_unit_sm100_structured_sparse_gemm_compressor_f32 + cutlass_test_unit_sm100_structured_sparse_gemm_compressor_f16 + cutlass_test_unit_sm100_structured_sparse_gemm_compressor_f8 + cutlass_test_unit_sm100_structured_sparse_gemm_compressor_f6 + cutlass_test_unit_sm100_structured_sparse_gemm_compressor_f4_qmma + cutlass_test_unit_sm100_structured_sparse_gemm_compressor_f4_omma +) + +cutlass_test_unit_add_executable( + cutlass_test_unit_sm100_structured_sparse_gemm_compressor_f32 + + sm100_sparse_gemm_compressor_f32.cu +) + +cutlass_test_unit_add_executable( + cutlass_test_unit_sm100_structured_sparse_gemm_compressor_f16 + + sm100_sparse_gemm_compressor_f16.cu +) + +cutlass_test_unit_add_executable( + cutlass_test_unit_sm100_structured_sparse_gemm_compressor_f8 + + sm100_sparse_gemm_compressor_f8.cu +) + +cutlass_test_unit_add_executable( + cutlass_test_unit_sm100_structured_sparse_gemm_compressor_f6 + + sm100_sparse_gemm_compressor_f6.cu +) + +cutlass_test_unit_add_executable( + cutlass_test_unit_sm100_structured_sparse_gemm_compressor_f4_qmma + + sm100_sparse_gemm_compressor_f4_qmma.cu +) + +cutlass_test_unit_add_executable( + cutlass_test_unit_sm100_structured_sparse_gemm_compressor_f4_omma + + sm100_sparse_gemm_compressor_f4_omma.cu +) + diff --git a/test/unit/transform/device/sm100_sparse_gemm_compressor_f16.cu b/test/unit/transform/device/sm100_sparse_gemm_compressor_f16.cu new file mode 100644 index 00000000..4e6d376c --- /dev/null +++ b/test/unit/transform/device/sm100_sparse_gemm_compressor_f16.cu @@ -0,0 +1,92 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +#include "cute/arch/mma_sm100_desc.hpp" // cute::UMMA::Major +#include "cutlass/gemm/collective/builders/sm100_common.inl" // tag_to_umma_major_A +#include "cutlass/gemm/collective/builders/sm1xx_sparse_config.inl" // Sm1xxGemmSparseConfig +#include "cutlass/transform/kernel/sparse_gemm_compressor.hpp" // StructuredSparseCompressor +#include "cutlass/transform/device/transform_universal_adapter.hpp" // TransformUniversalAdapter +#include "testbed_sparse_gemm_compressor.hpp" // TestbedSparseGemmCompressor + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// * Test Plan +// ElementA : fp16 +// LayoutA : row / col +// Gemm : 1x 2x 3x multiplier of alignment requirement. corner case that smaller than alignment requirement +/////////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f16_t) +{ + // Test Settings + using ElementA = cutlass::half_t; + using LayoutATag = cutlass::layout::RowMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementA>; + using ElementEMma = cute::sparse_elem<8, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto()); +} + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f16_n) +{ + // Test Settings + using ElementA = cutlass::bfloat16_t; + using LayoutATag = cutlass::layout::ColumnMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementA>; + using ElementEMma = cute::sparse_elem<8, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto() ); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/transform/device/sm100_sparse_gemm_compressor_f32.cu b/test/unit/transform/device/sm100_sparse_gemm_compressor_f32.cu new file mode 100644 index 00000000..985c35df --- /dev/null +++ b/test/unit/transform/device/sm100_sparse_gemm_compressor_f32.cu @@ -0,0 +1,92 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +#include "cute/arch/mma_sm100_desc.hpp" // cute::UMMA::Major +#include "cutlass/gemm/collective/builders/sm100_common.inl" // tag_to_umma_major_A +#include "cutlass/gemm/collective/builders/sm1xx_sparse_config.inl" // Sm1xxGemmSparseConfig +#include "cutlass/transform/kernel/sparse_gemm_compressor.hpp" // StructuredSparseCompressor +#include "cutlass/transform/device/transform_universal_adapter.hpp" // TransformUniversalAdapter +#include "testbed_sparse_gemm_compressor.hpp" // TestbedSparseGemmCompressor + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// * Test Plan +// ElementA : fp32 +// LayoutA : row / col +// Gemm : 1x 2x 3x multiplier of alignment requirement. corner case that smaller than alignment requirement +/////////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f32_t) +{ + // Test Settings + using ElementA = float; + using LayoutATag = cutlass::layout::RowMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementA>; + using ElementEMma = cute::sparse_elem<4, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto()); +} + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f32_n) +{ + // Test Settings + using ElementA = cutlass::tfloat32_t; + using LayoutATag = cutlass::layout::ColumnMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementA>; + using ElementEMma = cute::sparse_elem<4, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto()); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/transform/device/sm100_sparse_gemm_compressor_f4_omma.cu b/test/unit/transform/device/sm100_sparse_gemm_compressor_f4_omma.cu new file mode 100644 index 00000000..29cafdf3 --- /dev/null +++ b/test/unit/transform/device/sm100_sparse_gemm_compressor_f4_omma.cu @@ -0,0 +1,92 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +#include "cute/arch/mma_sm100_desc.hpp" // cute::UMMA::Major +#include "cutlass/gemm/collective/builders/sm100_common.inl" // tag_to_umma_major_A +#include "cutlass/gemm/collective/builders/sm1xx_sparse_config.inl" // Sm1xxGemmSparseConfig +#include "cutlass/transform/kernel/sparse_gemm_compressor.hpp" // StructuredSparseCompressor +#include "cutlass/transform/device/transform_universal_adapter.hpp" // TransformUniversalAdapter +#include "testbed_sparse_gemm_compressor.hpp" // TestbedSparseGemmCompressor + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// * Test Plan +// ElementA : fp4 (qmma), fp4 (omma) +// LayoutA : row / col +// Gemm : 1x 2x 3x multiplier of alignment requirement. corner case that smaller than alignment requirement +/////////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, omma_f4_t) +{ + // Test Settings + using ElementA = cutlass::float_e2m1_t; + using LayoutATag = cutlass::layout::RowMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<4, uint8_t>; + using ElementEMma = cute::sparse_elem<16, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto()); +} + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, omma_f4_runtimedtype_t) +{ + // Test Settings + using ElementA = cutlass::type_erased_dynamic_float4_t; + using LayoutATag = cutlass::layout::RowMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<4, uint8_t>; + using ElementEMma = cute::sparse_elem<16, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto_small()); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/transform/device/sm100_sparse_gemm_compressor_f4_qmma.cu b/test/unit/transform/device/sm100_sparse_gemm_compressor_f4_qmma.cu new file mode 100644 index 00000000..fb75ad77 --- /dev/null +++ b/test/unit/transform/device/sm100_sparse_gemm_compressor_f4_qmma.cu @@ -0,0 +1,139 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +#include "cute/arch/mma_sm100_desc.hpp" // cute::UMMA::Major +#include "cutlass/gemm/collective/builders/sm100_common.inl" // tag_to_umma_major_A +#include "cutlass/gemm/collective/builders/sm1xx_sparse_config.inl" // Sm1xxGemmSparseConfig +#include "cutlass/transform/kernel/sparse_gemm_compressor.hpp" // StructuredSparseCompressor +#include "cutlass/transform/device/transform_universal_adapter.hpp" // TransformUniversalAdapter +#include "testbed_sparse_gemm_compressor.hpp" // TestbedSparseGemmCompressor + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// * Test Plan +// ElementA : fp4 (qmma), fp4 (omma) +// LayoutA : row / col +// Gemm : 1x 2x 3x multiplier of alignment requirement. corner case that smaller than alignment requirement +/////////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f4_t) +{ + // Test Settings + using ElementA = cutlass::float_e2m1_t; + using LayoutATag = cutlass::layout::RowMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementA>; + using ElementEMma = cute::sparse_elem<8, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto()); +} + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f4_n) +{ + // Test Settings + using ElementA = cutlass::float_e2m1_t; + using LayoutATag = cutlass::layout::ColumnMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementA>; + using ElementEMma = cute::sparse_elem<8, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto()); +} + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f4_runtimedtype_t) +{ + // Test Settings + using ElementA = cutlass::type_erased_dynamic_float4_t; + using ElementAMmaRaw = cutlass::detail::type_erased_dynamic_float4_unpacksmem_t; + using LayoutATag = cutlass::layout::RowMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementAMmaRaw>; + using ElementEMma = cute::sparse_elem<8, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto_small()); +} + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f4_runtimedtype_n) +{ + // Test Settings + using ElementA = cutlass::type_erased_dynamic_float4_t; + using ElementAMmaRaw = cutlass::detail::type_erased_dynamic_float4_unpacksmem_t; + using LayoutATag = cutlass::layout::ColumnMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementAMmaRaw>; + using ElementEMma = cute::sparse_elem<8, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto_small()); +} + + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/transform/device/sm100_sparse_gemm_compressor_f6.cu b/test/unit/transform/device/sm100_sparse_gemm_compressor_f6.cu new file mode 100644 index 00000000..4be2a9ac --- /dev/null +++ b/test/unit/transform/device/sm100_sparse_gemm_compressor_f6.cu @@ -0,0 +1,138 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +#include "cute/arch/mma_sm100_desc.hpp" // cute::UMMA::Major +#include "cutlass/gemm/collective/builders/sm100_common.inl" // tag_to_umma_major_A +#include "cutlass/gemm/collective/builders/sm1xx_sparse_config.inl" // Sm1xxGemmSparseConfig +#include "cutlass/transform/kernel/sparse_gemm_compressor.hpp" // StructuredSparseCompressor +#include "cutlass/transform/device/transform_universal_adapter.hpp" // TransformUniversalAdapter +#include "testbed_sparse_gemm_compressor.hpp" // TestbedSparseGemmCompressor + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// * Test Plan +// ElementA : fp6 +// LayoutA : row / col +// Gemm : 1x 2x 3x multiplier of alignment requirement. corner case that smaller than alignment requirement +/////////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f6_t) +{ + // Test Settings + using ElementA = cutlass::float_e3m2_t; + using LayoutATag = cutlass::layout::RowMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementA>; + using ElementEMma = cute::sparse_elem<8, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto()); +} + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f6_n) +{ + // Test Settings + using ElementA = cutlass::float_e2m3_t; + using LayoutATag = cutlass::layout::ColumnMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementA>; + using ElementEMma = cute::sparse_elem<8, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto()); +} + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f6_runtimedtype_t) +{ + // Test Settings + using ElementA = cutlass::type_erased_dynamic_float6_t; + using ElementAMmaRaw = cutlass::detail::type_erased_dynamic_float6_unpacksmem_t; + using LayoutATag = cutlass::layout::RowMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementAMmaRaw>; + using ElementEMma = cute::sparse_elem<8, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto_small()); +} + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f6_runtimedtype_n) +{ + // Test Settings + using ElementA = cutlass::type_erased_dynamic_float6_t; + using ElementAMmaRaw = cutlass::detail::type_erased_dynamic_float6_unpacksmem_t; + using LayoutATag = cutlass::layout::ColumnMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementAMmaRaw>; + using ElementEMma = cute::sparse_elem<8, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto_small()); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/transform/device/sm100_sparse_gemm_compressor_f8.cu b/test/unit/transform/device/sm100_sparse_gemm_compressor_f8.cu new file mode 100644 index 00000000..e682ad0a --- /dev/null +++ b/test/unit/transform/device/sm100_sparse_gemm_compressor_f8.cu @@ -0,0 +1,136 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +#include "cute/arch/mma_sm100_desc.hpp" // cute::UMMA::Major +#include "cutlass/gemm/collective/builders/sm100_common.inl" // tag_to_umma_major_A +#include "cutlass/gemm/collective/builders/sm1xx_sparse_config.inl" // Sm1xxGemmSparseConfig +#include "cutlass/transform/kernel/sparse_gemm_compressor.hpp" // StructuredSparseCompressor +#include "cutlass/transform/device/transform_universal_adapter.hpp" // TransformUniversalAdapter +#include "testbed_sparse_gemm_compressor.hpp" // TestbedSparseGemmCompressor + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// * Test Plan +// ElementA : fp8 +// LayoutA : row / col +// Gemm : 1x 2x 3x multiplier of alignment requirement. corner case that smaller than alignment requirement +/////////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f8_t) +{ + // Test Settings + using ElementA = cutlass::float_e4m3_t; + using LayoutATag = cutlass::layout::RowMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementA>; + using ElementEMma = cute::sparse_elem<8, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto()); +} + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f8_n) +{ + // Test Settings + using ElementA = cutlass::float_e5m2_t; + using LayoutATag = cutlass::layout::ColumnMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementA>; + using ElementEMma = cute::sparse_elem<8, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto()); +} + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f8_runtimedtype_t) +{ + // Test Settings + using ElementA = cutlass::type_erased_dynamic_float8_t; + using LayoutATag = cutlass::layout::RowMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementA>; + using ElementEMma = cute::sparse_elem<8, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto_small()); +} + +TEST(SM100_Structured_Sparse_Gemm_Compressor_Device, f8_runtimedtype_n) +{ + // Test Settings + using ElementA = cutlass::type_erased_dynamic_float8_t; + using LayoutATag = cutlass::layout::ColumnMajor; + + // Deduct From Test Setting + using ElementAMma = cute::sparse_elem<2, ElementA>; + using ElementEMma = cute::sparse_elem<8, uint8_t>; + + using Sm1xxSparseConfig = cutlass::Sm1xxGemmSparseConfig; + + using CompressorKernel = cutlass::transform::kernel:: + StructuredSparseCompressor, ElementA, LayoutATag, Sm1xxSparseConfig, cutlass::arch::Sm100>; + + using Compressor = cutlass::transform::device::TransformUniversalAdapter; + + // Test Bed + test::transform::device::TestbedSparseGemmCompressor testbed; + EXPECT_TRUE(testbed.run_auto_small()); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/transform/device/testbed_sparse_gemm_compressor.hpp b/test/unit/transform/device/testbed_sparse_gemm_compressor.hpp index f4445824..a42505f2 100644 --- a/test/unit/transform/device/testbed_sparse_gemm_compressor.hpp +++ b/test/unit/transform/device/testbed_sparse_gemm_compressor.hpp @@ -107,6 +107,10 @@ initialize_tensor(cutlass::TensorView view, cutlass::Distributi scope_max = 2; scope_min = 0; } + else if (bits_input <= 6) { + scope_max = 2; + scope_min = -2; + } else if (bits_input <= 8) { scope_max = 1; scope_min = -1; @@ -155,9 +159,11 @@ public: using ElementA = typename CompressorKernel::ElementA; using LayoutATag = typename CompressorKernel::LayoutATag; using StrideA = typename CompressorKernel::StrideA; + static constexpr bool IsRuntimeDataTypeA = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4(); using ArrayElementA = - ElementA - ; + cute::conditional_t>, + ElementA>; using ElementE = typename CompressorKernel::ElementEMmaRaw; using LayoutETag = cutlass::layout::RowMajor; // We don't care about the major here, just to allocate tensor diff --git a/tools/library/include/cutlass/library/descriptions.h b/tools/library/include/cutlass/library/descriptions.h index 6f1dc5ff..152c61de 100644 --- a/tools/library/include/cutlass/library/descriptions.h +++ b/tools/library/include/cutlass/library/descriptions.h @@ -328,6 +328,7 @@ struct BlockScaleDescription { struct GroupedGemmDescription : public OperationDescription { GemmDescription gemm; std::optional block_scales; + bool is_moe{false}; }; /// Description of all GEMM computations diff --git a/tools/library/include/cutlass/library/library.h b/tools/library/include/cutlass/library/library.h index ca843ce1..4c2dbd32 100644 --- a/tools/library/include/cutlass/library/library.h +++ b/tools/library/include/cutlass/library/library.h @@ -592,6 +592,9 @@ struct GemmGroupedArguments { // underlying operation uses the one it needs. cute::Shape* problem_sizes_3x; cute::Shape* problem_sizes_3x_host; + std::vector max_problem_size_3x; + int32_t* tokens_per_expert{nullptr}; + int32_t* tokens_per_expert_host{nullptr}; }; struct GroupedGemmBlockScaledArguments : GemmGroupedArguments { diff --git a/tools/library/src/grouped_gemm_operation_3x.hpp b/tools/library/src/grouped_gemm_operation_3x.hpp index 91f618d4..fae564ca 100644 --- a/tools/library/src/grouped_gemm_operation_3x.hpp +++ b/tools/library/src/grouped_gemm_operation_3x.hpp @@ -869,5 +869,505 @@ public: } }; +template +class MoeGroupedGemmOperation3xBase : public GemmOperation3xBase { +public: + using Operator = Operator_; + using OperatorArguments = typename Operator::Arguments; + using ElementA = typename Operator::ElementA; + using LayoutA = typename Operator::LayoutA; + using ElementB = typename Operator::ElementB; + using LayoutB = typename Operator::LayoutB; + using ElementC = typename Operator::ElementC; + using LayoutC = typename Operator::LayoutC; + using ElementD = typename Operator::ElementD; + using LayoutD = typename Operator::LayoutD; + using ElementAccumulator = typename Operator::ElementAccumulator; + using ElementCompute = typename Operator::EpilogueOutputOp::ElementCompute; + using ArrayElementA = typename Operator::GemmKernel::CollectiveMainloop::ArrayElementA; + using ArrayElementB = typename Operator::GemmKernel::CollectiveMainloop::ArrayElementB; + + using CollectiveMainloop = typename Operator::CollectiveMainloop; + using CollectiveEpilogue = typename Operator::CollectiveEpilogue; + using ThreadEpilogueOp = typename CollectiveEpilogue::ThreadEpilogueOp; + using StrideC = typename Operator::GemmKernel::StrideC; + using StrideD = typename Operator::GemmKernel::StrideD; + + static constexpr bool IsRuntimeDataTypeA = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4(); + static constexpr bool IsRuntimeDataTypeB = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4(); + static_assert((IsRuntimeDataTypeA && IsRuntimeDataTypeB) || + (!IsRuntimeDataTypeA && !IsRuntimeDataTypeB), + "ElementA and ElementB in a GEMM kernel should be both runtime or both static."); + static constexpr bool IsRuntimeDataType = IsRuntimeDataTypeA && IsRuntimeDataTypeB; + + MoeGroupedGemmOperation3xBase(char const* name = "unknown_gemm") + : GemmOperation3xBase(name, GemmKind::kGrouped) { + this->description_.is_moe = true; + this->description_.kind = OperationKind::kGroupedGemm; + this->description_.name = name; + this->description_.provider = Provider::kCUTLASS; + + this->description_.gemm = GemmOperation3xBase::description_; + this->description_.tile_description = this->description_.gemm.tile_description; + }; + +public: + + // mutable CudaBuffer strideC_device; + // mutable CudaBuffer strideD_device; + + /// Returns the description of the GEMM operation + virtual OperationDescription const& description() const override final { return description_; } + /// Gets the host-side workspace + uint64_t get_host_workspace_size(void const* configuration) const override final { + return sizeof(Operator); + } + +protected: + library::GroupedGemmDescription description_; + + /// Constructs the arguments structure given the configuration and arguments + Status update_arguments_base( + OperatorArguments& operator_args, + GemmGroupedArguments const& arguments) const { + operator_args.mode = cutlass::gemm::GemmUniversalMode::kGrouped; + int M= arguments.max_problem_size_3x[0]; + int N = arguments.max_problem_size_3x[1]; + int K = arguments.max_problem_size_3x[2]; + int L = arguments.problem_count; + operator_args.problem_shape = { + M, + N, + K, + L, + arguments.tokens_per_expert, + arguments.tokens_per_expert_host + }; + + if constexpr (IsRuntimeDataType) { + using RuntimeDataTypeA = typename Operator::GemmKernel::CollectiveMainloop::RuntimeDataTypeA; + using RuntimeDataTypeB = typename Operator::GemmKernel::CollectiveMainloop::RuntimeDataTypeB; + + static_assert(cute::is_same_v, + "RuntimeDataTypeA/B should be identical, either MXF8F6F4Format or MXF4Format"); + using RuntimeDatatypeArg = RuntimeDataTypeA; + + auto mapping = [](RuntimeDatatype type) { + if constexpr (cute::is_same_v) { + if (type == RuntimeDatatype::kE5M2) { + return cute::UMMA::MXF8F6F4Format::E5M2; + } + else if (type == RuntimeDatatype::kE4M3) { + return cute::UMMA::MXF8F6F4Format::E4M3; + } + else if (type == RuntimeDatatype::kE3M2) { + return cute::UMMA::MXF8F6F4Format::E3M2; + } + else if (type == RuntimeDatatype::kE2M3) { + return cute::UMMA::MXF8F6F4Format::E2M3; + } + else if (type == RuntimeDatatype::kE2M1) { + return cute::UMMA::MXF8F6F4Format::E2M1; + } + else { + #if defined(CUTLASS_DEBUG_TRACE_LEVEL) && CUTLASS_DEBUG_TRACE_LEVEL >= 1 + std::cerr << "Invalid input datatype specified. Running with e4m3." << std::endl; + #endif + return cute::UMMA::MXF8F6F4Format::E4M3; + } + } + else if constexpr (cute::is_same_v) { + if (type == RuntimeDatatype::kE2M1) { + return cute::UMMA::MXF4Format::E2M1; + } + else { + #if defined(CUTLASS_DEBUG_TRACE_LEVEL) && CUTLASS_DEBUG_TRACE_LEVEL >= 1 + std::cerr << "Invalid input datatype specified. Running with e2m1." << std::endl; + #endif + return cute::UMMA::MXF4Format::E2M1; + } + } + // BlockScaled kernels receive either MXF4Format or MXF8F6F4Format runtime datatype + CUTE_GCC_UNREACHABLE; + }; + operator_args.mainloop.runtime_data_type_a = mapping(arguments.runtime_input_datatype_a); + operator_args.mainloop.runtime_data_type_b = mapping(arguments.runtime_input_datatype_b); + } + + operator_args.epilogue.ptr_C = static_cast(arguments.ptr_C); + operator_args.epilogue.ptr_D = static_cast(arguments.ptr_D); + + operator_args.epilogue.dC = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, L)); + operator_args.epilogue.dD = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, L)); + + /* Query device SM count and max active clusters to pass onto the kernel as an argument, where needed */ + operator_args.hw_info.sm_count = arguments.sm_count; + if constexpr (Operator::ArchTag::kMinComputeCapability >= 90) { + operator_args.hw_info.max_active_clusters = arguments.max_active_clusters; + } + if constexpr (!std::is_const_v) { + operator_args.scheduler.max_swizzle_size = arguments.swizzle_size; + } + + if constexpr (!std::is_const_v) { + using Enum_t = decltype(operator_args.scheduler.raster_order); + switch (arguments.raster_order) { + case RasterOrder::kAlongN: + operator_args.scheduler.raster_order = Enum_t::AlongN; + break; + case RasterOrder::kAlongM: + operator_args.scheduler.raster_order = Enum_t::AlongM; + break; + default: + operator_args.scheduler.raster_order = Enum_t::Heuristic; + } + } + + if constexpr (Operator::ArchTag::kMinComputeCapability >= 100) { + operator_args.hw_info.cluster_shape = + dim3(arguments.cluster_shape.m(), arguments.cluster_shape.n(), arguments.cluster_shape.k()); + operator_args.hw_info.cluster_shape_fallback = dim3( + arguments.cluster_shape_fallback.m(), + arguments.cluster_shape_fallback.n(), + arguments.cluster_shape_fallback.k()); + } + return Status::kSuccess; + } + + template + static Status update_fusion_args(FusionArgs& fusion_args, GemmGroupedArguments const& arguments) { + if (arguments.pointer_mode == ScalarPointerMode::kHost) { + fusion_args.alpha = *static_cast(arguments.alpha); + fusion_args.beta = *static_cast(arguments.beta); + fusion_args.alpha_ptr = nullptr; + fusion_args.beta_ptr = nullptr; + return Status::kSuccess; + } + else if (arguments.pointer_mode == ScalarPointerMode::kDevice) { + fusion_args.alpha = 0; + fusion_args.beta = 0; + fusion_args.alpha_ptr = static_cast(arguments.alpha); + fusion_args.beta_ptr = static_cast(arguments.beta); + return Status::kSuccess; + } + else { + return Status::kErrorInvalidProblem; + } + } +}; + +template +class MoeGroupedGemmUniversal3xOperation : public MoeGroupedGemmOperation3xBase { +public: + using Base = MoeGroupedGemmOperation3xBase; + using Operator = Operator_; + using OperatorArguments = typename Operator::Arguments; + + MoeGroupedGemmUniversal3xOperation(char const* name = "unknown_gemm") + : MoeGroupedGemmOperation3xBase(name) { + } + ~MoeGroupedGemmUniversal3xOperation() override = default; +protected: + template struct UpdateFusionArgs { + static Status update_(FusionArgs const& fusion_args, GemmGroupedArguments const& arguments) { + // If a custom EVT is instantiated then it is the users's responsibility + // to ensure alpha and beta are updated appropriately + return Status::kSuccess; + } + }; + + template + struct UpdateFusionArgs> { + static Status update_(FusionArgs& fusion_args, GemmGroupedArguments const& arguments) { + return MoeGroupedGemmOperation3xBase::update_fusion_args(fusion_args, arguments); + } + }; + + /// Constructs the arguments structure given the configuration and arguments + Status + update_arguments_(OperatorArguments& operator_args, GemmGroupedArguments const* arguments) const { + + Status status = UpdateFusionArgs::update_( + operator_args.epilogue.thread, + *arguments); + if (status != Status::kSuccess) { + return status; + } + + status = this->update_arguments_base(operator_args, *arguments); + + operator_args.mainloop.ptr_A = static_cast(arguments->ptr_A); + operator_args.mainloop.ptr_B = static_cast(arguments->ptr_B); + + return status; + } +public: + /// Returns success if the operation can proceed + Status can_implement([[maybe_unused]] void const* configuration_ptr, void const* arguments_ptr) + const override { + GemmGroupedArguments const* arguments = static_cast(arguments_ptr); + OperatorArguments args; + auto status = update_arguments_(args, arguments); + if (status != Status::kSuccess) { + return status; + } + + status = Operator::can_implement(args); + return status; + } + /// Gets the device-side workspace + uint64_t get_device_workspace_size(void const* configuration_ptr, void const* arguments_ptr) + const override { + + OperatorArguments args; + auto status = update_arguments_(args, static_cast(arguments_ptr)); + if (status != Status::kSuccess) { + return 0; + } + + uint64_t size = Operator::get_workspace_size(args); + return size; + } + /// Initializes the workspace + /// **** CAUTION **** + /// Must be called when ldc, or ldd change. + /// The CUTLASS library stores the operations in a type- + /// erased manifest. Therefore, only this class knows + /// the type of strideC, and strideD. + /// Since grouped GEMM needs to allocate storage for + /// the strides on device, the concrete type of the stride + /// must be known in order to copy in the correct memory + /// layout on device. + Status initialize( + void const* configuration_ptr, + void* host_workspace, + void* device_workspace, + cudaStream_t stream = nullptr) const override { + + Operator* op = new (host_workspace) Operator; + + return Status::kSuccess; + } + /// **** CAUTION **** + /// initialize() must be called if lda, ldb, ldc, or ldd change. + Status run( + void const* arguments_ptr, + void* host_workspace, + void* device_workspace = nullptr, + cudaStream_t stream = nullptr) const override { + + OperatorArguments operator_args; + auto const& args = *static_cast(arguments_ptr); + + Status status = update_arguments_(operator_args, &args); + if (status != Status::kSuccess) { + return status; + } + Operator* op = static_cast(host_workspace); + // We need to call initialize() since we have to rebuild TMA desc for every new set of args + status = op->run(operator_args, device_workspace, stream, nullptr, args.use_pdl); + return status; + } + // Set arguments that should only be set once before verifying or profiling the kernel. + // This should encompass any expensive operations that don't vary from run to run + // (e.g., max_active_clusters). + Status initialize_with_arguments(void* arguments_ptr) const override { + if constexpr (Operator::ArchTag::kMinComputeCapability < 90) { + return Status::kSuccess; + } + + GemmGroupedArguments* args = static_cast(arguments_ptr); + + dim3 cluster_dims; + if constexpr (cute::is_static_v) { + cluster_dims = dim3( + cute::size<0>(typename Operator::GemmKernel::ClusterShape{}), + cute::size<1>(typename Operator::GemmKernel::ClusterShape{}), + cute::size<2>(typename Operator::GemmKernel::ClusterShape{}) + ); + } + else { + cluster_dims = dim3( + args->cluster_shape.m(), + args->cluster_shape.n(), + args->cluster_shape.k() + ); + } + + uint32_t threads_per_block = Operator::GemmKernel::MaxThreadsPerBlock; + void const* kernel_ptr = (void*)(device_kernel); + args->max_active_clusters = cutlass::KernelHardwareInfo::query_device_max_active_clusters( + cluster_dims, + threads_per_block, + kernel_ptr); + + if (args->max_active_clusters == 0) { + std::cerr << "Max Active Clusters could not be queried. " + << "Falling back to heuristics mode (static cluster shape) or preferred cluster mode.\n"; + } + + return Status::kSuccess; + } +}; + +template +class BlockScaledMoeGroupedGemmUniversal3xOperation : public MoeGroupedGemmOperation3xBase { +public: + using Base = MoeGroupedGemmOperation3xBase; + using Operator = Operator_; + using OperatorArguments = typename Operator::Arguments; + using ElementD = typename Operator::ElementD; + using LayoutD = typename Operator::LayoutD; + using ElementAccumulator = typename Operator::ElementAccumulator; + using ElementCompute = typename Operator::EpilogueOutputOp::ElementCompute; + + using CollectiveMainloop = typename Operator::CollectiveMainloop; + using CollectiveEpilogue = typename Operator::CollectiveEpilogue; + using ThreadEpilogueOp = typename CollectiveEpilogue::ThreadEpilogueOp; + + using ElementSF = typename Operator::CollectiveMainloop::ElementSF; + + using TiledMma = typename Operator::CollectiveMainloop::TiledMma; + constexpr static int SFVecSize = TiledMma::SFVecSize; + + static constexpr bool epilogue_scalefactor_generation = not cute::is_same_v; + static constexpr int32_t SFD_VectorSize = epilogue_scalefactor_generation ? ThreadEpilogueOp::SFVecSize : SFVecSize; + using ElementSFD = cute::conditional_t; + using LayoutSFD = cute::conditional_t; + + BlockScaledMoeGroupedGemmUniversal3xOperation(char const* name = "unknown_gemm") + : MoeGroupedGemmOperation3xBase(name) { + BlockScaleDescription block_scaled_desc{}; + block_scaled_desc.kind = OperationKind::kBlockScaledGemm; + block_scaled_desc.SFA.element = NumericTypeMap::kId; + block_scaled_desc.SFA.layout = LayoutTypeID::kRowMajor; + block_scaled_desc.SFA.alignment = 128; + block_scaled_desc.SFA.log_extent_range = 32; + block_scaled_desc.SFA.log_stride_range = 32; + + block_scaled_desc.SFB.element = NumericTypeMap::kId; + block_scaled_desc.SFB.layout = LayoutTypeID::kRowMajor; + block_scaled_desc.SFB.alignment = 128; + block_scaled_desc.SFB.log_extent_range = 32; + block_scaled_desc.SFB.log_stride_range = 32; + + block_scaled_desc.SFMVecSize = 1; + block_scaled_desc.SFNVecSize = 1; + block_scaled_desc.SFKVecSize = SFVecSize; + + block_scaled_desc.SFD = make_TensorDescription(128); + block_scaled_desc.EpilogueSFVecSize = SFD_VectorSize; + + this->description_.block_scales = block_scaled_desc; + } + ~BlockScaledMoeGroupedGemmUniversal3xOperation() override = default; +protected: + template struct UpdateFusionArgs { + static Status update_(FusionArgs const& fusion_args, GroupedGemmBlockScaledArguments const& arguments) { + // If a custom EVT is instantiated then it is the users's responsibility + // to ensure alpha and beta are updated appropriately + return Status::kSuccess; + } + }; + template + struct UpdateFusionArgs> { + static Status + update_(FusionArgs& fusion_args, GroupedGemmBlockScaledArguments const& arguments) { + + if constexpr (epilogue_scalefactor_generation) { + fusion_args.block_scale_factor_ptr = static_cast(arguments.SFD); + fusion_args.norm_constant_ptr = static_cast(arguments.norm_constant); + } + + return MoeGroupedGemmOperation3xBase::update_fusion_args(fusion_args, arguments); + } + }; + +public: + /// Returns success if the operation can proceed + Status can_implement([[maybe_unused]] void const* configuration_ptr, void const* arguments_ptr) + const override { + GroupedGemmBlockScaledArguments const* arguments = + static_cast(arguments_ptr); + OperatorArguments args; + auto status = update_arguments_(args, arguments); + if (status != Status::kSuccess) { + return status; + } + + status = Operator::can_implement(args); + return status; + } + + Status update_arguments_( + OperatorArguments& operator_args, + GroupedGemmBlockScaledArguments const* arguments) const { + Status status = UpdateFusionArgs::update_( + operator_args.epilogue.thread, + *arguments); + if (status != Status::kSuccess) { + return status; + } + operator_args.mainloop.ptr_A = static_cast(arguments->ptr_A); + operator_args.mainloop.ptr_B = static_cast(arguments->ptr_B); + + operator_args.mainloop.ptr_SFA = static_cast(arguments->SFA); + operator_args.mainloop.ptr_SFB = static_cast(arguments->SFB); + return this->update_arguments_base(operator_args, *arguments); + } + + uint64_t get_device_workspace_size(void const* configuration_ptr, void const* arguments_ptr) + const override { + + OperatorArguments args; + auto status = + update_arguments_(args, static_cast(arguments_ptr)); + if (status != Status::kSuccess) { + return 0; + } + + uint64_t size = Operator::get_workspace_size(args); + return size; + } + + /// Initializes the workspace + /// **** CAUTION **** + /// Must be called when ldc, or ldd change. + /// The CUTLASS library stores the operations in a type- + /// erased manifest. Therefore, only this class knows + /// the type of strideA, strideB, strideC, and strideD. + /// Since grouped GEMM needs to allocate storage for + /// the strides on device, the concrete type of the stride + /// must be known in order to copy in the correct memory + /// layout on device. + Status initialize( + void const* configuration_ptr, + void* host_workspace, + void* device_workspace, + cudaStream_t stream = nullptr) const override { + Operator* op = new (host_workspace) Operator; + return Status::kSuccess; + } + /// **** CAUTION **** + /// initialize() must be called if ldc, or ldd change. + Status run( + void const* arguments_ptr, + void* host_workspace, + void* device_workspace = nullptr, + cudaStream_t stream = nullptr) const override { + + OperatorArguments operator_args; + auto const& args = *static_cast(arguments_ptr); + + Status status = update_arguments_(operator_args, &args); + if (status != Status::kSuccess) { + return status; + } + + Operator* op = static_cast(host_workspace); + status = op->run(operator_args, device_workspace, stream, nullptr); + return status; + } +}; } // namespace cutlass::library diff --git a/tools/library/src/handle.cu b/tools/library/src/handle.cu index 00ad2e0e..b1349125 100644 --- a/tools/library/src/handle.cu +++ b/tools/library/src/handle.cu @@ -279,6 +279,9 @@ static int gemm_problem_alignment( int max_element_alignment = 0; for (NumericTypeID type_id : elements) { + if (library::sizeof_bits(type_id)==0) { + continue; + } int element_alignment = max_alignment_in_bytes * 8 / library::sizeof_bits(type_id); max_element_alignment = std::max(max_element_alignment, element_alignment); } diff --git a/tools/library/src/reference/block_scaled_gemm_fp4a_vs16.cu b/tools/library/src/reference/block_scaled_gemm_fp4a_vs16.cu index 7afb6db5..89960db7 100644 --- a/tools/library/src/reference/block_scaled_gemm_fp4a_vs16.cu +++ b/tools/library/src/reference/block_scaled_gemm_fp4a_vs16.cu @@ -74,6 +74,11 @@ void initialize_block_scaled_gemm_reference_operations_fp4a_vs16(Manifest &manif 16 /*EpilogueSFVecSize*/ >(manifest); + make_block_scaled_gemm_tn< + float_e2m1_t /*A*/, float_ue4m3_t /*SFA*/, float_e2m1_t /*B*/, float_ue4m3_t /*SFB*/, + void /*C*/, float /*Compute*/, void /*SFD*/, float /*Accum*/, bfloat16_t /*D*/, 16 /*SFVecSize*/ + >(manifest); + make_block_scaled_gemm_tn< float_e2m1_t /*A*/, float_ue4m3_t /*SFA*/, float_e2m1_t /*B*/, float_ue4m3_t /*SFB*/, half_t /*C*/, float /*Compute*/, float_ue8m0_t /*SFD*/, float /*Accum*/, float_e2m1_t /*D*/, 16 /*SFVecSize*/, diff --git a/tools/library/src/reference/block_scaled_gemm_fp4a_vs32.cu b/tools/library/src/reference/block_scaled_gemm_fp4a_vs32.cu index b646d085..9cf20b99 100644 --- a/tools/library/src/reference/block_scaled_gemm_fp4a_vs32.cu +++ b/tools/library/src/reference/block_scaled_gemm_fp4a_vs32.cu @@ -73,6 +73,15 @@ void initialize_block_scaled_gemm_reference_operations_fp4a_vs32(Manifest &manif half_t /*C*/, float /*Compute*/, void /*SFD*/, float /*Accum*/, float_e3m2_t /*D*/, 32 /*SFVecSize*/ >(manifest); + make_block_scaled_gemm< + float_e2m1_t /*A*/, float_ue8m0_t /*SFA*/, float_e2m1_t /*B*/, float_ue8m0_t /*SFB*/, + bfloat16_t /*C*/, float /*Compute*/, void /*SFD*/, float /*Accum*/, bfloat16_t /*D*/, 32 /*SFVecSize*/ + >(manifest); + make_block_scaled_gemm< + float_e2m1_t /*A*/, float_ue8m0_t /*SFA*/, float_e2m1_t /*B*/, float_ue8m0_t /*SFB*/, + void /*C*/, float /*Compute*/, void /*SFD*/, float /*Accum*/, bfloat16_t /*D*/, 32 /*SFVecSize*/ + >(manifest); + // With SF generation reference make_block_scaled_gemm< float_e2m1_t /*A*/, float_ue8m0_t /*SFA*/, float_e2m1_t /*B*/, float_ue8m0_t /*SFB*/, diff --git a/tools/library/src/reference/gemm_e4m3a_e4m3out.cu b/tools/library/src/reference/gemm_e4m3a_e4m3out.cu index d45093e8..f140d364 100644 --- a/tools/library/src/reference/gemm_e4m3a_e4m3out.cu +++ b/tools/library/src/reference/gemm_e4m3a_e4m3out.cu @@ -56,6 +56,15 @@ void initialize_gemm_reference_operations_e4m3a_e4m3out(Manifest &manifest) { float_e4m3_t // ElementD >(manifest); + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + make_gemm_real_canonical_layouts< float_e4m3_t, // ElementA float_e5m2_t, // ElementB @@ -65,6 +74,15 @@ void initialize_gemm_reference_operations_e4m3a_e4m3out(Manifest &manifest) { float_e4m3_t // ElementD >(manifest); + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + make_gemm_real_canonical_layouts< float_e4m3_t, // ElementA float_e4m3_t, // ElementB diff --git a/tools/library/src/reference/gemm_e4m3a_e5m2out.cu b/tools/library/src/reference/gemm_e4m3a_e5m2out.cu index 9a444f96..8d935797 100644 --- a/tools/library/src/reference/gemm_e4m3a_e5m2out.cu +++ b/tools/library/src/reference/gemm_e4m3a_e5m2out.cu @@ -56,6 +56,15 @@ void initialize_gemm_reference_operations_e4m3a_e5m2out(Manifest &manifest) { float_e5m2_t // ElementD >(manifest); + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e5m2_t // ElementD + >(manifest); + make_gemm_real_canonical_layouts< float_e4m3_t, // ElementA float_e5m2_t, // ElementB @@ -65,6 +74,15 @@ void initialize_gemm_reference_operations_e4m3a_e5m2out(Manifest &manifest) { float_e5m2_t // ElementD >(manifest); + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e5m2_t // ElementD + >(manifest); + make_gemm_real_canonical_layouts< float_e4m3_t, // ElementA float_e4m3_t, // ElementB diff --git a/tools/library/src/reference/gemm_e5m2a_e4m3out.cu b/tools/library/src/reference/gemm_e5m2a_e4m3out.cu index e6c68e59..acf511ea 100644 --- a/tools/library/src/reference/gemm_e5m2a_e4m3out.cu +++ b/tools/library/src/reference/gemm_e5m2a_e4m3out.cu @@ -56,6 +56,15 @@ void initialize_gemm_reference_operations_e5m2a_e4m3out(Manifest &manifest) { float_e4m3_t // ElementD >(manifest); + make_gemm_real_canonical_layouts< + float_e5m2_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + make_gemm_real_canonical_layouts< float_e5m2_t, // ElementA float_e5m2_t, // ElementB @@ -65,6 +74,15 @@ void initialize_gemm_reference_operations_e5m2a_e4m3out(Manifest &manifest) { float_e4m3_t // ElementD >(manifest); + make_gemm_real_canonical_layouts< + float_e5m2_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + make_gemm_real_canonical_layouts< float_e5m2_t, // ElementA float_e4m3_t, // ElementB diff --git a/tools/library/src/reference/gemm_e5m2a_e5m2out.cu b/tools/library/src/reference/gemm_e5m2a_e5m2out.cu index 67247c5d..313aae0b 100644 --- a/tools/library/src/reference/gemm_e5m2a_e5m2out.cu +++ b/tools/library/src/reference/gemm_e5m2a_e5m2out.cu @@ -56,6 +56,15 @@ void initialize_gemm_reference_operations_e5m2a_e5m2out(Manifest &manifest) { float_e5m2_t // ElementD >(manifest); + make_gemm_real_canonical_layouts< + float_e5m2_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e5m2_t // ElementD + >(manifest); + make_gemm_real_canonical_layouts< float_e5m2_t, // ElementA float_e5m2_t, // ElementB diff --git a/tools/library/src/reference/gemm_f4_f4_f32.cu b/tools/library/src/reference/gemm_f4_f4_f32.cu index 7e1e67c4..be8e2176 100644 --- a/tools/library/src/reference/gemm_f4_f4_f32.cu +++ b/tools/library/src/reference/gemm_f4_f4_f32.cu @@ -97,7 +97,46 @@ void initialize_gemm_reference_operations_f4_f4_f32(Manifest &manifest) { float, // ElementAccumulator float // ElementD >(manifest); + + // 1. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e2m1_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + // 2. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e2m1_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e5m2_t // ElementD + >(manifest); + + // 3. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e2m1_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); + + // 4. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e2m1_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_f4_f6_f32.cu b/tools/library/src/reference/gemm_f4_f6_f32.cu index d5a48b72..ce24e21a 100644 --- a/tools/library/src/reference/gemm_f4_f6_f32.cu +++ b/tools/library/src/reference/gemm_f4_f6_f32.cu @@ -99,6 +99,45 @@ void initialize_gemm_reference_operations_f4_f6_f32(Manifest &manifest) { float // ElementD >(manifest); + // 1. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e3m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + + // 2. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e3m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e5m2_t // ElementD + >(manifest); + + // 3. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e3m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); + + // 4. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e3m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_f4_f8_f32.cu b/tools/library/src/reference/gemm_f4_f8_f32.cu index b57bd1ff..265f121d 100644 --- a/tools/library/src/reference/gemm_f4_f8_f32.cu +++ b/tools/library/src/reference/gemm_f4_f8_f32.cu @@ -99,6 +99,46 @@ void initialize_gemm_reference_operations_f4_f8_f32(Manifest &manifest) { float // ElementD >(manifest); + // 1. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + + // 2. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e5m2_t // ElementD + >(manifest); + + // 3. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); + + // 4. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); + // 1. make_gemm_real_canonical_layouts< float_e2m1_t, // ElementA @@ -139,6 +179,45 @@ void initialize_gemm_reference_operations_f4_f8_f32(Manifest &manifest) { float // ElementD >(manifest); + // 1. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + + // 2. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e5m2_t // ElementD + >(manifest); + + // 3. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); + + // 4. + make_gemm_real_canonical_layouts< + float_e2m1_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_f6_f4_f32.cu b/tools/library/src/reference/gemm_f6_f4_f32.cu index db03eaf8..b07032b6 100644 --- a/tools/library/src/reference/gemm_f6_f4_f32.cu +++ b/tools/library/src/reference/gemm_f6_f4_f32.cu @@ -99,6 +99,47 @@ void initialize_gemm_reference_operations_f6_f4_f32(Manifest &manifest) { float // ElementD >(manifest); + + // 1. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e2m1_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + + // 2. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e2m1_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e5m2_t // ElementD + >(manifest); + + // 3. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e2m1_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); + + // 4. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e2m1_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); + } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_f6_f6_f32.cu b/tools/library/src/reference/gemm_f6_f6_f32.cu index e090da9d..89ded966 100644 --- a/tools/library/src/reference/gemm_f6_f6_f32.cu +++ b/tools/library/src/reference/gemm_f6_f6_f32.cu @@ -98,6 +98,46 @@ void initialize_gemm_reference_operations_f6_f6_f32(Manifest &manifest) { float // ElementD >(manifest); + // 1. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e3m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + + // 2. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e3m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e5m2_t // ElementD + >(manifest); + + // 3. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e3m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); + + // 4. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e3m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); + } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_f6_f8_f32.cu b/tools/library/src/reference/gemm_f6_f8_f32.cu index cedeae2c..ee724453 100644 --- a/tools/library/src/reference/gemm_f6_f8_f32.cu +++ b/tools/library/src/reference/gemm_f6_f8_f32.cu @@ -99,6 +99,45 @@ void initialize_gemm_reference_operations_f6_f8_f32(Manifest &manifest) { float // ElementD >(manifest); + // 1. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + + // 2. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e5m2_t // ElementD + >(manifest); + + // 3. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); + + // 4. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_f8_f4_f32.cu b/tools/library/src/reference/gemm_f8_f4_f32.cu index 064cd226..ac5f64bd 100644 --- a/tools/library/src/reference/gemm_f8_f4_f32.cu +++ b/tools/library/src/reference/gemm_f8_f4_f32.cu @@ -99,6 +99,46 @@ void initialize_gemm_reference_operations_f8_f4_f32(Manifest &manifest) { float // ElementD >(manifest); + // 1. + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e2m1_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + + // 2. + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e2m1_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e5m2_t // ElementD + >(manifest); + + // 3. + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e2m1_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); + + // 4. + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e2m1_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); + } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_f8_f6_f32.cu b/tools/library/src/reference/gemm_f8_f6_f32.cu index 0a5bb9e9..24397a4d 100644 --- a/tools/library/src/reference/gemm_f8_f6_f32.cu +++ b/tools/library/src/reference/gemm_f8_f6_f32.cu @@ -99,6 +99,46 @@ void initialize_gemm_reference_operations_f8_f6_f32(Manifest &manifest) { float // ElementD >(manifest); + // 1. + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e3m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + + // 2. + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e3m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e5m2_t // ElementD + >(manifest); + + // 3. + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e3m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); + + // 4. + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e3m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); + } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_fp32out.cu b/tools/library/src/reference/gemm_fp32out.cu index 3c19b459..fbb3aa43 100644 --- a/tools/library/src/reference/gemm_fp32out.cu +++ b/tools/library/src/reference/gemm_fp32out.cu @@ -54,6 +54,15 @@ void initialize_gemm_reference_operations_fp32out(Manifest &manifest) { float // ElementAccumulator >(manifest); + make_gemm_real_canonical_layouts< + float, // ElementA + float, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float + >(manifest); + make_gemm_real_canonical_layouts< tfloat32_t, tfloat32_t, @@ -62,6 +71,15 @@ void initialize_gemm_reference_operations_fp32out(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + tfloat32_t, + tfloat32_t, + void, + float, + float, + float + >(manifest); + make_gemm_real_canonical_layouts< tfloat32_t, tfloat32_t, @@ -69,6 +87,14 @@ void initialize_gemm_reference_operations_fp32out(Manifest &manifest) { float, float >(manifest); + make_gemm_real_canonical_layouts< + tfloat32_t, + tfloat32_t, + void, + float, + float, + tfloat32_t + >(manifest); make_gemm_real_canonical_layouts< half_t, @@ -78,6 +104,15 @@ void initialize_gemm_reference_operations_fp32out(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + half_t, + half_t, + void, + float, + float, + half_t + >(manifest); + make_gemm_real_canonical_layouts< half_t, half_t, @@ -86,6 +121,14 @@ void initialize_gemm_reference_operations_fp32out(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + half_t, + half_t, + void, + float, + float, + float + >(manifest); make_gemm_real_canonical_layouts< bfloat16_t, bfloat16_t, @@ -94,6 +137,15 @@ void initialize_gemm_reference_operations_fp32out(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + bfloat16_t, + bfloat16_t, + void, + float, + float, + float + >(manifest); + make_gemm_real_canonical_layouts< bfloat16_t, bfloat16_t, @@ -101,6 +153,15 @@ void initialize_gemm_reference_operations_fp32out(Manifest &manifest) { float, float >(manifest); + + make_gemm_real_canonical_layouts< + bfloat16_t, + bfloat16_t, + void, + float, + float, + bfloat16_t + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_fp8in_bf16out.cu b/tools/library/src/reference/gemm_fp8in_bf16out.cu index 581def9b..a2a571f4 100644 --- a/tools/library/src/reference/gemm_fp8in_bf16out.cu +++ b/tools/library/src/reference/gemm_fp8in_bf16out.cu @@ -82,6 +82,42 @@ void initialize_gemm_reference_operations_fp8in_bf16out(Manifest &manifest) { float, // ElementAccumulator bfloat16_t // ElementD >(manifest); + + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + bfloat16_t // ElementD + >(manifest); + + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + bfloat16_t // ElementD + >(manifest); + + make_gemm_real_canonical_layouts< + float_e5m2_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + bfloat16_t // ElementD + >(manifest); + + make_gemm_real_canonical_layouts< + float_e5m2_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + bfloat16_t // ElementD + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_fp8in_fp16out.cu b/tools/library/src/reference/gemm_fp8in_fp16out.cu index 8d8178ef..5bdfc1e8 100644 --- a/tools/library/src/reference/gemm_fp8in_fp16out.cu +++ b/tools/library/src/reference/gemm_fp8in_fp16out.cu @@ -82,6 +82,42 @@ void initialize_gemm_reference_operations_fp8in_fp16out(Manifest &manifest) { float, // ElementAccumulator half_t // ElementD >(manifest); + + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); + + make_gemm_real_canonical_layouts< + float_e5m2_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); + + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); + + make_gemm_real_canonical_layouts< + float_e5m2_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_fp8in_fp32out.cu b/tools/library/src/reference/gemm_fp8in_fp32out.cu index 005ba8be..2827ac93 100644 --- a/tools/library/src/reference/gemm_fp8in_fp32out.cu +++ b/tools/library/src/reference/gemm_fp8in_fp32out.cu @@ -82,6 +82,42 @@ void initialize_gemm_reference_operations_fp8in_fp32out(Manifest &manifest) { float, // ElementAccumulator float // ElementD >(manifest); + + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); + + make_gemm_real_canonical_layouts< + float_e4m3_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); + + make_gemm_real_canonical_layouts< + float_e5m2_t, // ElementA + float_e4m3_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); + + make_gemm_real_canonical_layouts< + float_e5m2_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_fp_mixed_input.cu b/tools/library/src/reference/gemm_fp_mixed_input.cu index 18fbf2ff..a4bdc6ea 100644 --- a/tools/library/src/reference/gemm_fp_mixed_input.cu +++ b/tools/library/src/reference/gemm_fp_mixed_input.cu @@ -54,6 +54,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + int8_t, + half_t, + void, + float, + float, + float + >(manifest); + make_gemm_real_canonical_layouts< uint8_t, half_t, @@ -61,6 +70,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + uint8_t, + half_t, + void, + float, + float, + float + >(manifest); + make_gemm_real_canonical_layouts< int8_t, half_t, @@ -68,6 +86,16 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + int8_t, + half_t, + void, + float, + float, + half_t + >(manifest); + + make_gemm_real_canonical_layouts< uint8_t, half_t, @@ -75,6 +103,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + uint8_t, + half_t, + void, + float, + float, + half_t + >(manifest); + make_gemm_real_canonical_layouts< int8_t, half_t, @@ -82,6 +119,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { half_t >(manifest); + make_gemm_real_canonical_layouts< + int8_t, + half_t, + void, + half_t, + half_t, + half_t + >(manifest); + make_gemm_real_canonical_layouts< uint8_t, half_t, @@ -89,6 +135,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { half_t >(manifest); + make_gemm_real_canonical_layouts< + uint8_t, + half_t, + void, + half_t, + half_t, + half_t + >(manifest); + make_gemm_real_canonical_layouts< half_t, int8_t, @@ -96,6 +151,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + half_t, + int8_t, + void, + float, + float, + float + >(manifest); + make_gemm_real_canonical_layouts< half_t, uint8_t, @@ -103,6 +167,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + half_t, + uint8_t, + void, + float, + float, + float + >(manifest); + make_gemm_real_canonical_layouts< half_t, int8_t, @@ -110,6 +183,16 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { half_t >(manifest); + + make_gemm_real_canonical_layouts< + half_t, + int8_t, + void, + half_t, + half_t, + half_t + >(manifest); + make_gemm_real_canonical_layouts< half_t, uint8_t, @@ -117,6 +200,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { half_t >(manifest); + make_gemm_real_canonical_layouts< + half_t, + uint8_t, + void, + half_t, + half_t, + half_t + >(manifest); + make_gemm_real_canonical_layouts< half_t, int8_t, @@ -124,6 +216,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + half_t, + int8_t, + void, + float, + float, + half_t + >(manifest); + make_gemm_real_canonical_layouts< half_t, uint8_t, @@ -131,6 +232,16 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + + make_gemm_real_canonical_layouts< + half_t, + uint8_t, + void, + float, + float, + half_t + >(manifest); + // bfloat16_t mixed with 8-bit integer input make_gemm_real_canonical_layouts< int8_t, @@ -139,6 +250,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + int8_t, + bfloat16_t, + void, + float, + float, + float + >(manifest); + make_gemm_real_canonical_layouts< uint8_t, bfloat16_t, @@ -146,6 +266,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + uint8_t, + bfloat16_t, + void, + float, + float, + float + >(manifest); + make_gemm_real_canonical_layouts< int8_t, bfloat16_t, @@ -153,6 +282,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + int8_t, + bfloat16_t, + void, + float, + float, + bfloat16_t + >(manifest); + make_gemm_real_canonical_layouts< uint8_t, bfloat16_t, @@ -160,6 +298,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + uint8_t, + bfloat16_t, + void, + float, + float, + bfloat16_t + >(manifest); + make_gemm_real_canonical_layouts< bfloat16_t, int8_t, @@ -167,6 +314,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + bfloat16_t, + int8_t, + void, + float, + float, + float + >(manifest); + make_gemm_real_canonical_layouts< bfloat16_t, uint8_t, @@ -174,6 +330,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + bfloat16_t, + uint8_t, + void, + float, + float, + float + >(manifest); + make_gemm_real_canonical_layouts< bfloat16_t, int8_t, @@ -181,12 +346,30 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) { float >(manifest); + make_gemm_real_canonical_layouts< + bfloat16_t, + int8_t, + void, + float, + float, + bfloat16_t + >(manifest); + make_gemm_real_canonical_layouts< bfloat16_t, uint8_t, bfloat16_t, float >(manifest); + + make_gemm_real_canonical_layouts< + bfloat16_t, + uint8_t, + void, + float, + float, + bfloat16_t + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_fp_other.cu b/tools/library/src/reference/gemm_fp_other.cu index 5e9ad777..373b3925 100644 --- a/tools/library/src/reference/gemm_fp_other.cu +++ b/tools/library/src/reference/gemm_fp_other.cu @@ -54,6 +54,15 @@ void initialize_gemm_reference_operations_fp_other(Manifest &manifest) { half_t >(manifest); + make_gemm_real_canonical_layouts< + half_t, + half_t, + void, + half_t, + half_t, + half_t + >(manifest); + make_gemm_real_canonical_layouts< half_t, half_t, @@ -62,6 +71,24 @@ void initialize_gemm_reference_operations_fp_other(Manifest &manifest) { half_t >(manifest); + make_gemm_real_canonical_layouts< + half_t, + half_t, + float, + float, + float, + half_t + >(manifest); + + make_gemm_real_canonical_layouts< + bfloat16_t, + bfloat16_t, + float, + float, + float, + bfloat16_t + >(manifest); + make_gemm_real_canonical_layouts< double, double, @@ -70,6 +97,15 @@ void initialize_gemm_reference_operations_fp_other(Manifest &manifest) { double >(manifest); + make_gemm_real_canonical_layouts< + double, + double, + void, + double, + double, + double + >(manifest); + make_gemm_complex_canonical_layouts< complex, complex, @@ -78,6 +114,15 @@ void initialize_gemm_reference_operations_fp_other(Manifest &manifest) { complex >(manifest); + make_gemm_complex_canonical_layouts< + complex, + complex, + void, + complex, + complex, + complex + >(manifest); + make_gemm_complex_canonical_layouts< complex, complex, @@ -85,6 +130,15 @@ void initialize_gemm_reference_operations_fp_other(Manifest &manifest) { complex, complex >(manifest); + + make_gemm_complex_canonical_layouts< + complex, + complex, + void, + complex, + complex, + complex + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_int4.cu b/tools/library/src/reference/gemm_int4.cu index ffaf7d34..a12eca5f 100644 --- a/tools/library/src/reference/gemm_int4.cu +++ b/tools/library/src/reference/gemm_int4.cu @@ -59,7 +59,28 @@ void initialize_gemm_reference_operations_int4(Manifest &manifest) { 64, int4b_t, int4b_t, + void, int32_t, + int32_t, + int32_t + >(manifest); + + make_gemm_interleaved_layouts< + 64, + int4b_t, + int4b_t, + int32_t, + float, + int32_t, + int32_t, + NumericConverterClamp + >(manifest); + + make_gemm_interleaved_layouts< + 64, + int4b_t, + int4b_t, + void, float, int32_t, int32_t, @@ -77,6 +98,17 @@ void initialize_gemm_reference_operations_int4(Manifest &manifest) { NumericConverterClamp >(manifest); + make_gemm_interleaved_layouts< + 64, + int4b_t, + int4b_t, + void, + float, + int32_t, + int4b_t, + NumericConverterClamp + >(manifest); + make_gemm_interleaved_layouts< 64, uint4b_t, @@ -90,7 +122,28 @@ void initialize_gemm_reference_operations_int4(Manifest &manifest) { 64, uint4b_t, uint4b_t, + void, int32_t, + int32_t, + int32_t + >(manifest); + + make_gemm_interleaved_layouts< + 64, + uint4b_t, + uint4b_t, + int32_t, + float, + int32_t, + int32_t, + NumericConverterClamp + >(manifest); + + make_gemm_interleaved_layouts< + 64, + uint4b_t, + uint4b_t, + void, float, int32_t, int32_t, @@ -108,6 +161,17 @@ void initialize_gemm_reference_operations_int4(Manifest &manifest) { NumericConverterClamp >(manifest); + make_gemm_interleaved_layouts< + 64, + uint4b_t, + uint4b_t, + void, + float, + int32_t, + uint4b_t, + NumericConverterClamp + >(manifest); + make_gemm_interleaved_layouts< 64, uint4b_t, @@ -118,6 +182,17 @@ void initialize_gemm_reference_operations_int4(Manifest &manifest) { int4b_t, NumericConverterClamp >(manifest); + + make_gemm_interleaved_layouts< + 64, + uint4b_t, + uint4b_t, + void, + float, + int32_t, + int4b_t, + NumericConverterClamp + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_int8_interleaved_32.cu b/tools/library/src/reference/gemm_int8_interleaved_32.cu index 7885afc0..874e19f9 100644 --- a/tools/library/src/reference/gemm_int8_interleaved_32.cu +++ b/tools/library/src/reference/gemm_int8_interleaved_32.cu @@ -59,7 +59,28 @@ void initialize_gemm_reference_operations_int8_interleaved_32(Manifest &manifest 32, int8_t, int8_t, + void, int32_t, + int32_t, + int32_t + >(manifest); + + make_gemm_interleaved_layouts< + 32, + int8_t, + int8_t, + int32_t, + float, + int32_t, + int32_t, + NumericConverterClamp + >(manifest); + + make_gemm_interleaved_layouts< + 32, + int8_t, + int8_t, + void, float, int32_t, int32_t, @@ -77,6 +98,17 @@ void initialize_gemm_reference_operations_int8_interleaved_32(Manifest &manifest NumericConverterClamp >(manifest); + make_gemm_interleaved_layouts< + 32, + int8_t, + int8_t, + void, + float, + int32_t, + int8_t, + NumericConverterClamp + >(manifest); + make_gemm_interleaved_layouts< 32, uint8_t, @@ -90,7 +122,28 @@ void initialize_gemm_reference_operations_int8_interleaved_32(Manifest &manifest 32, uint8_t, uint8_t, + void, int32_t, + int32_t, + int32_t + >(manifest); + + make_gemm_interleaved_layouts< + 32, + uint8_t, + uint8_t, + int32_t, + float, + int32_t, + int32_t, + NumericConverterClamp + >(manifest); + + make_gemm_interleaved_layouts< + 32, + uint8_t, + uint8_t, + void, float, int32_t, int32_t, @@ -108,6 +161,17 @@ void initialize_gemm_reference_operations_int8_interleaved_32(Manifest &manifest NumericConverterClamp >(manifest); + make_gemm_interleaved_layouts< + 32, + uint8_t, + uint8_t, + void, + float, + int32_t, + uint8_t, + NumericConverterClamp + >(manifest); + make_gemm_interleaved_layouts< 32, uint8_t, @@ -118,6 +182,17 @@ void initialize_gemm_reference_operations_int8_interleaved_32(Manifest &manifest int8_t, NumericConverterClamp >(manifest); + + make_gemm_interleaved_layouts< + 32, + uint8_t, + uint8_t, + void, + float, + int32_t, + int8_t, + NumericConverterClamp + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_int8_interleaved_64.cu b/tools/library/src/reference/gemm_int8_interleaved_64.cu index effc47ac..11d751a0 100644 --- a/tools/library/src/reference/gemm_int8_interleaved_64.cu +++ b/tools/library/src/reference/gemm_int8_interleaved_64.cu @@ -59,7 +59,28 @@ void initialize_gemm_reference_operations_int8_interleaved_64(Manifest &manifest 64, int4b_t, int4b_t, + void, int32_t, + int32_t, + int32_t + >(manifest); + + make_gemm_interleaved_layouts< + 64, + int4b_t, + int4b_t, + int32_t, + float, + int32_t, + int32_t, + NumericConverterClamp + >(manifest); + + make_gemm_interleaved_layouts< + 64, + int4b_t, + int4b_t, + void, float, int32_t, int32_t, @@ -77,6 +98,18 @@ void initialize_gemm_reference_operations_int8_interleaved_64(Manifest &manifest NumericConverterClamp >(manifest); + make_gemm_interleaved_layouts< + 64, + int4b_t, + int4b_t, + void, + float, + int32_t, + int4b_t, + NumericConverterClamp + >(manifest); + + make_gemm_interleaved_layouts< 64, uint4b_t, @@ -90,7 +123,28 @@ void initialize_gemm_reference_operations_int8_interleaved_64(Manifest &manifest 64, uint4b_t, uint4b_t, + void, int32_t, + int32_t, + int32_t + >(manifest); + + make_gemm_interleaved_layouts< + 64, + uint4b_t, + uint4b_t, + int32_t, + float, + int32_t, + int32_t, + NumericConverterClamp + >(manifest); + + make_gemm_interleaved_layouts< + 64, + uint4b_t, + uint4b_t, + void, float, int32_t, int32_t, @@ -108,6 +162,17 @@ void initialize_gemm_reference_operations_int8_interleaved_64(Manifest &manifest NumericConverterClamp >(manifest); + make_gemm_interleaved_layouts< + 64, + uint4b_t, + uint4b_t, + void, + float, + int32_t, + uint4b_t, + NumericConverterClamp + >(manifest); + make_gemm_interleaved_layouts< 64, uint4b_t, @@ -118,6 +183,17 @@ void initialize_gemm_reference_operations_int8_interleaved_64(Manifest &manifest int4b_t, NumericConverterClamp >(manifest); + + make_gemm_interleaved_layouts< + 64, + uint4b_t, + uint4b_t, + void, + float, + int32_t, + int4b_t, + NumericConverterClamp + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_int_mixed_input.cu b/tools/library/src/reference/gemm_int_mixed_input.cu index 0bf07ceb..1e42e29f 100644 --- a/tools/library/src/reference/gemm_int_mixed_input.cu +++ b/tools/library/src/reference/gemm_int_mixed_input.cu @@ -57,7 +57,26 @@ void initialize_gemm_reference_operations_int_mixed_input(Manifest &manifest) { make_gemm_real_canonical_layouts< int4b_t, int8_t, + void, + int32_t, + int32_t, + int32_t + >(manifest); + + make_gemm_real_canonical_layouts< + int4b_t, int8_t, + int8_t, + int32_t, + int32_t, + int8_t, + NumericConverterClamp + >(manifest); + + make_gemm_real_canonical_layouts< + int4b_t, + int8_t, + void, int32_t, int32_t, int8_t, @@ -77,7 +96,27 @@ void initialize_gemm_reference_operations_int_mixed_input(Manifest &manifest) { make_gemm_real_canonical_layouts< int4b_t, int8_t, + void, + float, + int32_t, + int32_t, + NumericConverterClamp + >(manifest); + + make_gemm_real_canonical_layouts< + int4b_t, int8_t, + int8_t, + float, + int32_t, + int8_t, + NumericConverterClamp + >(manifest); + + make_gemm_real_canonical_layouts< + int4b_t, + int8_t, + void, float, int32_t, int8_t, @@ -91,6 +130,15 @@ void initialize_gemm_reference_operations_int_mixed_input(Manifest &manifest) { int32_t >(manifest); + make_gemm_real_canonical_layouts< + int8_t, + int4b_t, + void, + int32_t, + int32_t, + int32_t + >(manifest); + make_gemm_real_canonical_layouts< int8_t, int4b_t, @@ -104,7 +152,27 @@ void initialize_gemm_reference_operations_int_mixed_input(Manifest &manifest) { make_gemm_real_canonical_layouts< int8_t, int4b_t, + void, int32_t, + int32_t, + int8_t, + NumericConverterClamp + >(manifest); + + make_gemm_real_canonical_layouts< + int8_t, + int4b_t, + int32_t, + float, + int32_t, + int32_t, + NumericConverterClamp + >(manifest); + + make_gemm_real_canonical_layouts< + int8_t, + int4b_t, + void, float, int32_t, int32_t, @@ -120,6 +188,16 @@ void initialize_gemm_reference_operations_int_mixed_input(Manifest &manifest) { int8_t, NumericConverterClamp >(manifest); + + make_gemm_real_canonical_layouts< + int8_t, + int4b_t, + void, + float, + int32_t, + int8_t, + NumericConverterClamp + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_reference_operation.h b/tools/library/src/reference/gemm_reference_operation.h index e07158b0..c65dd49d 100644 --- a/tools/library/src/reference/gemm_reference_operation.h +++ b/tools/library/src/reference/gemm_reference_operation.h @@ -86,7 +86,8 @@ public: using ElementC = ElementC_; using LayoutC = LayoutC_; using ElementD = ElementD_; - using TensorRefC = TensorRef; + using NonVoidElementC = cute::conditional_t, ElementD, ElementC>; + using TensorRefC = TensorRef; using TensorRefD = TensorRef; using ElementCompute = ElementCompute_; using ElementAccumulator = ElementAccumulator_; @@ -199,7 +200,7 @@ public: TensorRefA ref_A{static_cast(const_cast(args.A)), LayoutA(int(config.lda))}; TensorRefB ref_B{static_cast(const_cast(args.B)), LayoutB(int(config.ldb))}; - TensorRefC ref_C{static_cast(const_cast(args.C)), LayoutC(int(config.ldc))}; + TensorRefC ref_C{static_cast(const_cast(args.C)), LayoutC(int(config.ldc))}; TensorRefD ref_D{static_cast(args.D), LayoutC(int(config.ldd))}; if (kProvider == Provider::kReferenceHost) { @@ -209,7 +210,7 @@ public: LayoutA, ElementB, LayoutB, - ElementC, + NonVoidElementC, LayoutC, ElementCompute, ElementAccumulator, @@ -243,7 +244,7 @@ public: LayoutA, ElementB, LayoutB, - ElementC, + NonVoidElementC, LayoutC, ElementCompute, ElementAccumulator, @@ -434,7 +435,7 @@ template < typename ElementCompute_, typename ElementAccumulator_ = ElementCompute_, typename ElementD_ = ElementC_, - typename ConvertOp_ = NumericConverter, + typename ConvertOp_ = NumericConverter, typename InnerProductOp_ = multiply_add > void make_gemm_interleaved_layouts(Manifest &manifest) { diff --git a/tools/library/src/reference/gemm_s8_s8_s32.cu b/tools/library/src/reference/gemm_s8_s8_s32.cu index 939450f7..6594f286 100644 --- a/tools/library/src/reference/gemm_s8_s8_s32.cu +++ b/tools/library/src/reference/gemm_s8_s8_s32.cu @@ -135,6 +135,80 @@ void initialize_gemm_reference_operations_s8_s8_s32(Manifest &manifest) { half_t, // ElementD NumericConverterClamp // From Scalar to D >(manifest); + + // 1. + make_gemm_real_canonical_layouts< + int8_t, // ElementA + int8_t, // ElementB + void, // ElementC + int32_t, // ElementScalar / ElementCompute + int32_t, // ElementAccumulator + int32_t // ElementD + >(manifest); + + // 2. + make_gemm_real_canonical_layouts< + int8_t, // ElementA + int8_t, // ElementB + void, // ElementC + float, // ElementScalar / ElementCompute + int32_t, // ElementAccumulator + int32_t // ElementD + >(manifest); + + // 3. + make_gemm_real_canonical_layouts< + int8_t, // ElementA + int8_t, // ElementB + void, // ElementC + float, // ElementScalar / ElementCompute + int32_t, // ElementAccumulator + int8_t, // ElementD + NumericConverterClamp // From Scalar to D + >(manifest); + + // 4. + make_gemm_real_canonical_layouts< + int8_t, // ElementA + int8_t, // ElementB + void, // ElementC + int32_t, // ElementScalar / ElementCompute + int32_t, // ElementAccumulator + int8_t, // ElementD + NumericConverterClamp // From Scalar to D + >(manifest); + + // 5. + make_gemm_real_canonical_layouts< + int8_t, // ElementA + int8_t, // ElementB + void, // ElementC + float, // ElementScalar / ElementCompute + int32_t, // ElementAccumulator + int8_t, // ElementD + NumericConverterClamp // From Scalar to D + >(manifest); + + // 6. + make_gemm_real_canonical_layouts< + int8_t, // ElementA + int8_t, // ElementB + void, // ElementC + float, // ElementScalar / ElementCompute + int32_t, // ElementAccumulator + float // ElementD + >(manifest); + + // 7. + make_gemm_real_canonical_layouts< + int8_t, // ElementA + int8_t, // ElementB + void, // ElementC + float, // ElementScalar / ElementCompute + int32_t, // ElementAccumulator + half_t, // ElementD + NumericConverterClamp // From Scalar to D + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/library/src/reference/gemm_u8_u8_s32.cu b/tools/library/src/reference/gemm_u8_u8_s32.cu index 1f2786eb..30bf68e8 100644 --- a/tools/library/src/reference/gemm_u8_u8_s32.cu +++ b/tools/library/src/reference/gemm_u8_u8_s32.cu @@ -98,6 +98,49 @@ void initialize_gemm_reference_operations_u8_u8_s32(Manifest &manifest) { NumericConverterClamp // From Scalar to D >(manifest); + // 1. + make_gemm_real_canonical_layouts< + uint8_t, // ElementA + uint8_t, // ElementB + void, // ElementC + int32_t, // ElementScalar / ElementCompute + int32_t, // ElementAccumulator + int32_t // ElementD + >(manifest); + + // 2. + make_gemm_real_canonical_layouts< + uint8_t, // ElementA + uint8_t, // ElementB + void, // ElementC + float, // ElementScalar / ElementCompute + int32_t, // ElementAccumulator + int32_t, // ElementD + NumericConverterClamp // From Scalar to D + >(manifest); + + // 3. + make_gemm_real_canonical_layouts< + uint8_t, // ElementA + uint8_t, // ElementB + void, // ElementC + float, // ElementScalar / ElementCompute + int32_t, // ElementAccumulator + int8_t, // ElementD + NumericConverterClamp // From Scalar to D + >(manifest); + + // 4. + make_gemm_real_canonical_layouts< + uint8_t, // ElementA + uint8_t, // ElementB + void, // ElementC + float, // ElementScalar / ElementCompute + int32_t, // ElementAccumulator + uint8_t, // ElementD + NumericConverterClamp // From Scalar to D + >(manifest); + } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/profiler/include/cutlass/profiler/device_allocation.h b/tools/profiler/include/cutlass/profiler/device_allocation.h index 488b635c..cba847b3 100644 --- a/tools/profiler/include/cutlass/profiler/device_allocation.h +++ b/tools/profiler/include/cutlass/profiler/device_allocation.h @@ -84,6 +84,8 @@ private: /// The device ID where the allocation is made int device_; + /// Whether to free the memory when the object is destroyed + bool free_memory_{true}; public: // // Static member functions @@ -140,6 +142,16 @@ public: int batch_count = 1, int device = -1); + DeviceAllocation( + library::NumericTypeID type, + library::LayoutTypeID layout_id, + std::vector const &extent, + std::vector const &stride, + void* ref_pointer_, + int batch_count, + int device + ); + ~DeviceAllocation(); DeviceAllocation &reset(); diff --git a/tools/profiler/include/cutlass/profiler/device_context.h b/tools/profiler/include/cutlass/profiler/device_context.h index 0443b340..80981481 100644 --- a/tools/profiler/include/cutlass/profiler/device_context.h +++ b/tools/profiler/include/cutlass/profiler/device_context.h @@ -90,6 +90,18 @@ public: int batch_count, size_t device_index); +/// creates a reference tensor of existing memory, a given type, capacity (elements), and name + DeviceAllocation *create_ref_tensor( + Options const &options, + std::string const &name, + library::NumericTypeID type, + library::LayoutTypeID layout_id, + std::vector const &extent, + std::vector const &stride, + void* ref_pointer_, + int batch_count, + size_t device_index); + /// Allocates memory of a given type, capacity (elements), and name DeviceAllocation *allocate_and_initialize_tensor( Options const &options, diff --git a/tools/profiler/include/cutlass/profiler/grouped_gemm_operation_profiler.h b/tools/profiler/include/cutlass/profiler/grouped_gemm_operation_profiler.h index 62d47990..bea6d31f 100644 --- a/tools/profiler/include/cutlass/profiler/grouped_gemm_operation_profiler.h +++ b/tools/profiler/include/cutlass/profiler/grouped_gemm_operation_profiler.h @@ -68,6 +68,7 @@ public: std::vector problem_sizes; std::vector> problem_sizes_3x; + std::vector max_problem_size_3x; /// For exploration purposes std::vector> preferred_clusters; @@ -85,10 +86,13 @@ public: std::vector lda{0}; std::vector ldb{0}; std::vector ldc{0}; - + int64_t max_lda{0}; + int64_t max_ldb{0}; + int64_t max_ldc{0}; std::vector alpha; std::vector beta; + cutlass::library::RasterOrder raster_order{cutlass::library::RasterOrder::kHeuristic}; int swizzle_size{1}; @@ -168,7 +172,8 @@ public: DeviceAllocation* ldb_array_device{nullptr}; DeviceAllocation* ldc_array_device{nullptr}; DeviceAllocation* ldd_array_device{nullptr}; - + std::vector tokens_per_expert_host; + DeviceAllocation* tokens_per_expert_device{nullptr}; std::optional block_scales; library::GemmGroupedConfiguration configuration; @@ -188,6 +193,10 @@ private: arguments.ptr_B = gemm_workspace_.B_ptr_array_device[0]->data(); arguments.ptr_C = gemm_workspace_.C_ptr_array_device[0]->data(); arguments.ptr_D = gemm_workspace_.D_ptr_array_device[0]->data(); + if (is_moe) { + arguments.tokens_per_expert_host = gemm_workspace_.tokens_per_expert_host.data(); + arguments.tokens_per_expert = static_cast(gemm_workspace_.tokens_per_expert_device->data()); + } arguments.alpha = problem_.alpha.data(); arguments.beta = problem_.beta.data(); @@ -200,10 +209,11 @@ private: static_cast(gemm_workspace_.problem_sizes_array_device->data()); arguments.problem_sizes_3x = static_cast*>( gemm_workspace_.problem_sizes_3x_array_device->data()); - gemm_workspace_.arguments.problem_sizes_3x_host = problem_.problem_sizes_3x.data(); - gemm_workspace_.arguments.problem_count = problem_.problem_sizes.size(); - gemm_workspace_.arguments.cluster_shape = {int(problem_.cluster_m), int(problem_.cluster_n), int(problem_.cluster_k)}; - gemm_workspace_.arguments.cluster_shape_fallback = {int(problem_.cluster_m_fallback), int(problem_.cluster_n_fallback), int(problem_.cluster_k_fallback)}; + arguments.problem_sizes_3x_host = problem_.problem_sizes_3x.data(); + arguments.max_problem_size_3x = problem_.max_problem_size_3x; + arguments.problem_count = problem_.problem_sizes.size(); + arguments.cluster_shape = {int(problem_.cluster_m), int(problem_.cluster_n), int(problem_.cluster_k)}; + arguments.cluster_shape_fallback = {int(problem_.cluster_m_fallback), int(problem_.cluster_n_fallback), int(problem_.cluster_k_fallback)}; /* Query device SM count to pass onto the kernel as an argument, where needed */ arguments.sm_count = options.device.get_sm_count(0); @@ -230,7 +240,7 @@ protected: bool is_block_scaled{false}; bool is_blockwise{false}; - + bool is_moe{false}; public: GroupedGemmOperationProfiler(Options const& options); diff --git a/tools/profiler/src/device_allocation.cu b/tools/profiler/src/device_allocation.cu index 26efacdc..0f27abfd 100644 --- a/tools/profiler/src/device_allocation.cu +++ b/tools/profiler/src/device_allocation.cu @@ -326,8 +326,38 @@ DeviceAllocation::DeviceAllocation( reset(type, layout_id, extent, stride, batch_count); } +DeviceAllocation::DeviceAllocation( + library::NumericTypeID type, + library::LayoutTypeID layout_id, + std::vector const &extent, + std::vector const &stride, + void* ref_pointer_, + int batch_count, + int device +): + type_(type), batch_stride_(size_t(0)), capacity_(size_t(0)), + pointer_(ref_pointer_), batch_count_(1), device_(device), free_memory_(false) { + + tensor_ref_buffer_.resize(sizeof(pointer_) + (sizeof(int64_t) * library::get_layout_stride_rank(layout_id)), 0); + + type_ = type; + + layout_ = layout_id; + stride_ = stride; + extent_ = extent; + batch_count_ = batch_count; + + batch_stride_ = construct_layout( + tensor_ref_buffer_.data() + sizeof(pointer_), + layout_id, + extent, + stride_); + + capacity_ = batch_stride_ * batch_count_; +} + DeviceAllocation::~DeviceAllocation() { - if (pointer_) { + if (pointer_ and free_memory_) { int current_device; cudaGetDevice(¤t_device); @@ -343,7 +373,7 @@ DeviceAllocation::~DeviceAllocation() { } DeviceAllocation &DeviceAllocation::reset() { - if (pointer_) { + if (pointer_ and free_memory_) { int current_device; cudaGetDevice(¤t_device); @@ -366,6 +396,7 @@ DeviceAllocation &DeviceAllocation::reset() { extent_.clear(); tensor_ref_buffer_.clear(); batch_count_ = 1; + free_memory_ = true; return *this; } diff --git a/tools/profiler/src/device_context.cu b/tools/profiler/src/device_context.cu index 629770a3..ff649a6f 100644 --- a/tools/profiler/src/device_context.cu +++ b/tools/profiler/src/device_context.cu @@ -75,6 +75,27 @@ DeviceAllocation *DeviceContext::allocate_tensor( return allocation; } +/// creates a reference tensor of existing ptr, a given type, capacity (elements), and name +DeviceAllocation *DeviceContext::create_ref_tensor( + Options const &options, + std::string const &name, + library::NumericTypeID type, + library::LayoutTypeID layout_id, + std::vector const &extent, + std::vector const &stride, + void* ref_pointer_, + int batch_count, + size_t device_index) { + + int device = options.device.device_id(device_index); + device_memory_.emplace_back(type, layout_id, extent, stride, ref_pointer_, batch_count, + device); + DeviceAllocation *allocation = &device_memory_.back(); + + allocations_[name] = allocation; + return allocation; +} + static void initialize_allocation_with_data_distribution( Options const &options, int seed_shift, diff --git a/tools/profiler/src/gemm_operation_profiler.cu b/tools/profiler/src/gemm_operation_profiler.cu index 48e451e7..c97fd63f 100644 --- a/tools/profiler/src/gemm_operation_profiler.cu +++ b/tools/profiler/src/gemm_operation_profiler.cu @@ -1264,7 +1264,7 @@ bool GemmOperationProfiler::verify_cutlass( } } - // if verification.required is set, then return success iff at least one ref-check was run + // if verification.required is set, then return success if at least one ref-check was run if (options.verification.required) { bool did_any_verification_run = false; for (auto provider : options.verification.providers) { diff --git a/tools/profiler/src/grouped_gemm_operation_profiler.cu b/tools/profiler/src/grouped_gemm_operation_profiler.cu index 2297fab1..8ff36291 100644 --- a/tools/profiler/src/grouped_gemm_operation_profiler.cu +++ b/tools/profiler/src/grouped_gemm_operation_profiler.cu @@ -177,7 +177,7 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse( library::GroupedGemmDescription const& operation_desc, ProblemSpace const& problem_space, ProblemSpace::Problem const& problem) { - + bool is_moe = operation_desc.is_moe; this->mode = library::GemmUniversalMode::kGrouped; std::bitset<3> args_exist; @@ -189,7 +189,7 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse( arg_as_int(k, "k", problem_space, problem); std::string problem_file; args_exist[2] = arg_as_string(problem_file, "problem-sizes-file", problem_space, problem); - + int max_m = 0, max_n = 0, max_k = 0; if (args_exist.count() == 0) { int num_groups = 8; problem_sizes.resize(num_groups); @@ -204,6 +204,7 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse( problem_sizes[i] = {m, n, k}; problem_sizes_3x[i] = {m, n, k}; } + max_problem_size_3x = {m0 * num_groups, n0 * num_groups, k0 * num_groups}; } else if (args_exist.count() > 1) { std::cerr @@ -220,9 +221,13 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse( auto m = problems[i][0]; auto n = problems[i][1]; auto k = problems[i][2]; + max_m = std::max(max_m, m); + max_n = std::max(max_n, n); + max_k = std::max(max_k, k); problem_sizes[i] = {m, n, k}; problem_sizes_3x[i] = {m, n, k}; } + max_problem_size_3x = {max_m, max_n, max_k}; } // m, n, k path else if (args_exist[1]) { @@ -237,6 +242,7 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse( problem_sizes[i] = {m, n, k}; problem_sizes_3x[i] = {m, n, k}; } + max_problem_size_3x = {m, n, k}; } // --problem-sizes-file path else if (args_exist[2]) { @@ -247,7 +253,6 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse( // clear the problem sizes and 3x problem sizes from previous operation problem_sizes.clear(); problem_sizes_3x.clear(); - for (std::string line; std::getline(file, line);) { std::istringstream iss(line); @@ -258,14 +263,27 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse( if (iss >> m >> sep1 >> n >> sep2 >> k && sep1 == 'x' && sep2 == 'x' && !(iss >> remaining)) { problem_sizes.emplace_back(m, n, k); problem_sizes_3x.emplace_back(m, n, k); + max_m = std::max(max_m, m); + max_n = std::max(max_n, n); + max_k = std::max(max_k, k); } else { throw std::runtime_error( "Invalid format in line: " + line + ". Each line in file expected to be 'mxnxk'."); } } + max_problem_size_3x = {max_m, max_n, max_k}; + } + if (is_moe) { + for(size_t group_idx = 0; group_idx < problem_sizes.size(); group_idx++) { + if (problem_sizes[group_idx].m() != max_problem_size_3x[0] || + problem_sizes[group_idx].k() != max_problem_size_3x[2]) { + std::cerr << "Problem size M:"<< problem_sizes[group_idx].m() << "K:" << problem_sizes[group_idx].k() << " for group " << group_idx << "should be equal to " + << "Max problem size M:" << max_problem_size_3x[0] << "K:" << max_problem_size_3x[2] << " in MoE Grouped GEMM" << std::endl; + return Status::kErrorInvalidProblem; + } + } } - if (!arg_as_int(this->cluster_m, "cluster_m", problem_space, problem)) { // default value this->cluster_m = std::string(operation_desc.gemm.name).find("_2sm") != std::string::npos ? 2 : 1; @@ -382,6 +400,9 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse( operation_desc.gemm.C.layout, {int(this->m(group_idx)), int(this->n(group_idx))}) .front(); + this->max_lda = std::max(this->max_lda, this->lda[group_idx]); + this->max_ldb = std::max(this->max_ldb, this->ldb[group_idx]); + this->max_ldc = std::max(this->max_ldc, this->ldc[group_idx]); } // instantiation for exploration profiling @@ -609,7 +630,7 @@ Status GroupedGemmOperationProfiler::initialize_configuration( is_block_scaled = false; gemm_workspace_.block_scales = std::nullopt; } - + is_moe = operation_desc.is_moe; if (operation_desc.gemm.gemm_kind != library::GemmKind::kGrouped) { return Status::kErrorInvalidProblem; } @@ -761,232 +782,561 @@ Status GroupedGemmOperationProfiler::initialize_workspace( gemm_workspace_.reference_ptr_array_host.resize(num_groups); int seed_shift = 0; - for (size_t group_idx = 0; group_idx < num_groups; group_idx++) { - auto group_str = std::to_string(group_idx); - gemm_workspace_.A_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor( + if (not operation_desc.is_moe) { + for (size_t group_idx = 0; group_idx < num_groups; group_idx++) { + auto group_str = std::to_string(group_idx); + gemm_workspace_.A_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor( + options, + "A_" + group_str, + operation_desc.gemm.A.element, + operation_desc.gemm.A.layout, + {int(problem_.m(group_idx)), int(problem_.k(group_idx))}, + {int(problem_.lda[group_idx])}, + gemm_workspace_.problem_count, + seed_shift++, + 0); + gemm_workspace_.B_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor( + options, + "B_" + group_str, + operation_desc.gemm.B.element, + operation_desc.gemm.B.layout, + {int(problem_.k(group_idx)), int(problem_.n(group_idx))}, + {int(problem_.ldb[group_idx])}, + gemm_workspace_.problem_count, + seed_shift++, + 0); + gemm_workspace_.C_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor( + options, + "C_" + group_str, + operation_desc.gemm.C.element, + operation_desc.gemm.C.layout, + {int(problem_.m(group_idx)), int(problem_.n(group_idx))}, + {int(problem_.ldc[group_idx])}, + gemm_workspace_.problem_count, + seed_shift++, + 0); + gemm_workspace_.D_ptr_array_host[group_idx] = device_context.allocate_tensor( + options, + "D_" + group_str, + operation_desc.gemm.D.element, + operation_desc.gemm.D.layout, + {int(problem_.m(group_idx)), int(problem_.n(group_idx))}, + {int(problem_.ldc[group_idx])}, + gemm_workspace_.problem_count, + 0); + + gemm_workspace_.reference_ptr_array_host[group_idx] = device_context.allocate_tensor( + options, + "Reference_" + group_str, + operation_desc.gemm.D.element, + operation_desc.gemm.D.layout, + {int(problem_.m(group_idx)), int(problem_.n(group_idx))}, + {int(problem_.ldc[group_idx])}, + 1, + 0); + + if (is_block_scaled) { + auto const block_scale_desc = operation_desc.block_scales.value(); + auto& block_scale_ws = gemm_workspace_.block_scales.value(); + int sfa_m = round_up(int(problem_.m(group_idx)), 128); + int sfb_n = round_up(int(problem_.n(group_idx)), 128); + int sfa_sfb_k = + round_up(ceil_div(int(problem_.k(group_idx)), block_scale_desc.SFKVecSize), 4); + + int sfd_m = + block_scale_desc.SFD.layout == cutlass::library::LayoutTypeID::kRowMajor + ? sfa_m + : round_up(ceil_div(int(problem_.m(group_idx)), block_scale_desc.EpilogueSFVecSize), 4); + int sfd_n = + block_scale_desc.SFD.layout == cutlass::library::LayoutTypeID::kRowMajor + ? round_up(ceil_div(int(problem_.n(group_idx)), block_scale_desc.EpilogueSFVecSize), 4) + : sfb_n; + + block_scale_ws.SFA_ptr_array_host[group_idx] = + device_context.allocate_and_initialize_tensor( + options, + "SFA", + block_scale_desc.SFA.element, + block_scale_desc.SFA.layout, + {sfa_m, sfa_sfb_k}, + {sfa_sfb_k}, + gemm_workspace_.problem_count, + seed_shift++, + 0); + + block_scale_ws.SFB_ptr_array_host[group_idx] = + device_context.allocate_and_initialize_tensor( + options, + "SFB", + block_scale_desc.SFB.element, + block_scale_desc.SFB.layout, + {sfb_n, sfa_sfb_k}, + {sfa_sfb_k}, + gemm_workspace_.problem_count, + seed_shift++, + 0); + + block_scale_ws.SFD_ptr_array_host[group_idx] = device_context.allocate_tensor( + options, + "SFD", + block_scale_desc.SFD.element, + block_scale_desc.SFD.layout, + {sfd_m, sfd_n}, + {sfd_n}, + gemm_workspace_.problem_count, + 0); + + block_scale_ws.SFD_reference_ptr_array_host[group_idx] = device_context.allocate_tensor( + options, + "Reference_SFD", + block_scale_desc.SFD.element, + block_scale_desc.SFD.layout, + {sfd_m, sfd_n}, + {sfd_n}, + gemm_workspace_.problem_count, + 0); + + // ScaleFactor tensor results may have some holes and will not be touched by the kernel. + // If we randomly fill the two tensors, these holes may encounter refcheck errors. + if (block_scale_ws.SFD_ptr_array_host[group_idx]->type() != library::NumericTypeID::kVoid) { + block_scale_ws.SFD_reference_ptr_array_host[group_idx]->fill_device(0); + block_scale_ws.SFD_ptr_array_host[group_idx]->fill_device(0); + } + } + else if (is_blockwise) { + auto const block_scale_desc = operation_desc.block_scales.value(); + auto& block_scale_ws = gemm_workspace_.block_scales.value(); + int sfa_m = ceil_div(int(problem_.m(group_idx)), block_scale_desc.SFMVecSize); + int sfb_n = ceil_div(int(problem_.n(group_idx)), block_scale_desc.SFNVecSize); + int sfa_sfb_k = ceil_div(int(problem_.k(group_idx)), block_scale_desc.SFKVecSize); + + block_scale_ws.SFA_ptr_array_host[group_idx] = + device_context.allocate_and_initialize_tensor( + options, + "SFA_" + std::to_string(group_idx), + block_scale_desc.SFA.element, + block_scale_desc.SFA.layout, + {sfa_m, sfa_sfb_k}, + {sfa_m}, + gemm_workspace_.problem_count, + seed_shift++, + 0); + + block_scale_ws.SFB_ptr_array_host[group_idx] = + device_context.allocate_and_initialize_tensor( + options, + "SFB_" + std::to_string(group_idx), + block_scale_desc.SFB.element, + block_scale_desc.SFB.layout, + {sfa_sfb_k, sfb_n}, + {sfb_n}, + gemm_workspace_.problem_count, + seed_shift++, + 0); + } + } + + // takes the allocated tensors and initializes an array of pointers per problem in the workspace + auto create_dev_ptr_array_all_workspace = [&]( + std::vector& dev_ptr_arrays, + std::vector const& input, + std::string const& id) { + auto num_workspaces = gemm_workspace_.problem_count; + dev_ptr_arrays.resize(num_workspaces); + // note "problem_count" here refers to input/output count for L2 cycling + for (int i = 0; i < gemm_workspace_.problem_count; i++) { + std::string name = id + "_ptr_array_workspace" + std::to_string(i); + dev_ptr_arrays[i] = + device_context.allocate_block(options, name, library::NumericTypeID::kU64, num_groups, 0); + std::vector group_ptrs(num_groups); + for (size_t group_idx = 0; group_idx < num_groups; group_idx++) { + group_ptrs[group_idx] = input[group_idx]->batch_data(i); + } + dev_ptr_arrays[i]->copy_from_host(group_ptrs.data()); + } + }; + create_dev_ptr_array_all_workspace( + gemm_workspace_.A_ptr_array_device, + gemm_workspace_.A_ptr_array_host, + "A"); + create_dev_ptr_array_all_workspace( + gemm_workspace_.B_ptr_array_device, + gemm_workspace_.B_ptr_array_host, + "B"); + create_dev_ptr_array_all_workspace( + gemm_workspace_.C_ptr_array_device, + gemm_workspace_.C_ptr_array_host, + "C"); + create_dev_ptr_array_all_workspace( + gemm_workspace_.D_ptr_array_device, + gemm_workspace_.D_ptr_array_host, + "D"); + + if (is_block_scaled) { + auto& block_scale_ws = gemm_workspace_.block_scales.value(); + create_dev_ptr_array_all_workspace( + block_scale_ws.SFA_ptr_array_device, + block_scale_ws.SFA_ptr_array_host, + "SFA"); + create_dev_ptr_array_all_workspace( + block_scale_ws.SFB_ptr_array_device, + block_scale_ws.SFB_ptr_array_host, + "SFB"); + create_dev_ptr_array_all_workspace( + block_scale_ws.SFD_ptr_array_device, + block_scale_ws.SFD_ptr_array_host, + "SFD"); + + block_scale_ws.norm_constant = device_context.allocate_and_initialize_tensor( + options, + "norm_constant", + operation_desc.gemm.element_epilogue, + operation_desc.gemm.A.layout, // copied, but should this be D layout? + {1, 1}, + {1}, + 1, + seed_shift++, + 0 // device_index + ); + } + else if (is_blockwise) { + auto& block_scale_ws = gemm_workspace_.block_scales.value(); + create_dev_ptr_array_all_workspace( + block_scale_ws.SFA_ptr_array_device, + block_scale_ws.SFA_ptr_array_host, + "SFA"); + create_dev_ptr_array_all_workspace( + block_scale_ws.SFB_ptr_array_device, + block_scale_ws.SFB_ptr_array_host, + "SFB"); + } + } else { + int max_m = problem_.max_problem_size_3x[0]; + int max_n = problem_.max_problem_size_3x[1]; + int max_k = problem_.max_problem_size_3x[2]; + // allocate the block tensors + DeviceAllocation* block_A; + DeviceAllocation* block_B; + DeviceAllocation* block_C; + DeviceAllocation* block_D; + DeviceAllocation* block_ref_D; + DeviceAllocation* block_ref_SFD; + DeviceAllocation* block_SFA; + DeviceAllocation* block_SFB; + DeviceAllocation* block_SFD; + + gemm_workspace_.tokens_per_expert_host.resize(num_groups); + gemm_workspace_.tokens_per_expert_device = device_context.allocate_block( options, - "A_" + group_str, + "tokens_per_expert", + library::NumericTypeID::kU32, + num_groups, + 0); + block_A = device_context.allocate_and_initialize_tensor( + options, + "block_A", operation_desc.gemm.A.element, operation_desc.gemm.A.layout, - {int(problem_.m(group_idx)), int(problem_.k(group_idx))}, - {int(problem_.lda[group_idx])}, - gemm_workspace_.problem_count, + {max_m, max_k}, + {int(problem_.max_lda)}, + gemm_workspace_.problem_count * num_groups, seed_shift++, 0); - gemm_workspace_.B_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor( + block_B = device_context.allocate_and_initialize_tensor( options, - "B_" + group_str, + "block_B", operation_desc.gemm.B.element, operation_desc.gemm.B.layout, - {int(problem_.k(group_idx)), int(problem_.n(group_idx))}, - {int(problem_.ldb[group_idx])}, - gemm_workspace_.problem_count, + {max_k, max_n}, + {int(problem_.max_ldb)}, + gemm_workspace_.problem_count * num_groups, seed_shift++, 0); - gemm_workspace_.C_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor( + block_C = device_context.allocate_and_initialize_tensor( options, - "C_" + group_str, + "block_C", operation_desc.gemm.C.element, operation_desc.gemm.C.layout, - {int(problem_.m(group_idx)), int(problem_.n(group_idx))}, - {int(problem_.ldc[group_idx])}, - gemm_workspace_.problem_count, + {max_m, max_n}, + {int(problem_.max_ldc)}, + gemm_workspace_.problem_count * num_groups, seed_shift++, 0); - gemm_workspace_.D_ptr_array_host[group_idx] = device_context.allocate_tensor( + block_D = device_context.allocate_tensor( options, - "D_" + group_str, + "block_D", operation_desc.gemm.D.element, operation_desc.gemm.D.layout, - {int(problem_.m(group_idx)), int(problem_.n(group_idx))}, - {int(problem_.ldc[group_idx])}, - gemm_workspace_.problem_count, + {max_m, max_n}, + {int(problem_.max_ldc)}, + gemm_workspace_.problem_count * num_groups, + 0); + block_ref_D = device_context.allocate_tensor( + options, + "Block_Reference", + operation_desc.gemm.D.element, + operation_desc.gemm.D.layout, + {max_m, max_n}, + {int(problem_.max_ldc)}, + num_groups, 0); - gemm_workspace_.reference_ptr_array_host[group_idx] = device_context.allocate_tensor( - options, - "Reference_" + group_str, - operation_desc.gemm.D.element, - operation_desc.gemm.D.layout, - {int(problem_.m(group_idx)), int(problem_.n(group_idx))}, - {int(problem_.ldc[group_idx])}, - 1, - 0); if (is_block_scaled) { auto const block_scale_desc = operation_desc.block_scales.value(); auto& block_scale_ws = gemm_workspace_.block_scales.value(); - int sfa_m = round_up(int(problem_.m(group_idx)), 128); - int sfb_n = round_up(int(problem_.n(group_idx)), 128); + int sfa_m = round_up(problem_.max_problem_size_3x[0], 128); + int sfb_n = round_up(max_n, 128); int sfa_sfb_k = - round_up(ceil_div(int(problem_.k(group_idx)), block_scale_desc.SFKVecSize), 4); + round_up(ceil_div(max_k, block_scale_desc.SFKVecSize), 4); int sfd_m = block_scale_desc.SFD.layout == cutlass::library::LayoutTypeID::kRowMajor ? sfa_m - : round_up(ceil_div(int(problem_.m(group_idx)), block_scale_desc.EpilogueSFVecSize), 4); + : round_up(ceil_div(max_m, block_scale_desc.EpilogueSFVecSize), 4); int sfd_n = block_scale_desc.SFD.layout == cutlass::library::LayoutTypeID::kRowMajor - ? round_up(ceil_div(int(problem_.n(group_idx)), block_scale_desc.EpilogueSFVecSize), 4) + ? round_up(ceil_div(max_n, block_scale_desc.EpilogueSFVecSize), 4) : sfb_n; - - block_scale_ws.SFA_ptr_array_host[group_idx] = + block_SFA = device_context.allocate_and_initialize_tensor( options, - "SFA", + "block_SFA", block_scale_desc.SFA.element, block_scale_desc.SFA.layout, {sfa_m, sfa_sfb_k}, {sfa_sfb_k}, - gemm_workspace_.problem_count, + gemm_workspace_.problem_count * num_groups, seed_shift++, 0); - - block_scale_ws.SFB_ptr_array_host[group_idx] = + block_SFB = device_context.allocate_and_initialize_tensor( options, - "SFB", + "block_SFB", block_scale_desc.SFB.element, block_scale_desc.SFB.layout, {sfb_n, sfa_sfb_k}, {sfa_sfb_k}, - gemm_workspace_.problem_count, + gemm_workspace_.problem_count * num_groups, seed_shift++, 0); - block_scale_ws.SFD_ptr_array_host[group_idx] = device_context.allocate_tensor( + block_SFD = device_context.allocate_tensor( options, - "SFD", + "block_SFD", block_scale_desc.SFD.element, block_scale_desc.SFD.layout, {sfd_m, sfd_n}, {sfd_n}, - gemm_workspace_.problem_count, + gemm_workspace_.problem_count * num_groups, 0); - - block_scale_ws.SFD_reference_ptr_array_host[group_idx] = device_context.allocate_tensor( + block_ref_SFD = device_context.allocate_tensor( options, - "Reference_SFD", + "block_Reference_SFD", block_scale_desc.SFD.element, block_scale_desc.SFD.layout, {sfd_m, sfd_n}, {sfd_n}, - gemm_workspace_.problem_count, + gemm_workspace_.problem_count * num_groups, 0); - - // ScaleFactor tensor results may have some holes and will not be touched by the kernel. - // If we randomly fill the two tensors, these holes may encounter refcheck errors. - if (block_scale_ws.SFD_ptr_array_host[group_idx]->type() != library::NumericTypeID::kVoid) { - block_scale_ws.SFD_reference_ptr_array_host[group_idx]->fill_device(0); - block_scale_ws.SFD_ptr_array_host[group_idx]->fill_device(0); - } - } - else if (is_blockwise) { - auto const block_scale_desc = operation_desc.block_scales.value(); - auto& block_scale_ws = gemm_workspace_.block_scales.value(); - int sfa_m = ceil_div(int(problem_.m(group_idx)), block_scale_desc.SFMVecSize); - int sfb_n = ceil_div(int(problem_.n(group_idx)), block_scale_desc.SFNVecSize); - int sfa_sfb_k = ceil_div(int(problem_.k(group_idx)), block_scale_desc.SFKVecSize); - - block_scale_ws.SFA_ptr_array_host[group_idx] = - device_context.allocate_and_initialize_tensor( - options, - "SFA_" + std::to_string(group_idx), - block_scale_desc.SFA.element, - block_scale_desc.SFA.layout, - {sfa_m, sfa_sfb_k}, - {sfa_m}, - gemm_workspace_.problem_count, - seed_shift++, - 0); - - block_scale_ws.SFB_ptr_array_host[group_idx] = - device_context.allocate_and_initialize_tensor( - options, - "SFB_" + std::to_string(group_idx), - block_scale_desc.SFB.element, - block_scale_desc.SFB.layout, - {sfa_sfb_k, sfb_n}, - {sfb_n}, - gemm_workspace_.problem_count, - seed_shift++, - 0); - } - } - - // takes the allocated tensors and initializes an array of pointers per problem in the workspace - auto create_dev_ptr_array_all_workspace = [&]( - std::vector& dev_ptr_arrays, - std::vector const& input, - std::string const& id) { - auto num_workspaces = gemm_workspace_.problem_count; - dev_ptr_arrays.resize(num_workspaces); - // note "problem_count" here refers to input/output count for L2 cycling - for (int i = 0; i < gemm_workspace_.problem_count; i++) { - std::string name = id + "_ptr_array_workspace" + std::to_string(i); - dev_ptr_arrays[i] = - device_context.allocate_block(options, name, library::NumericTypeID::kU64, num_groups, 0); - std::vector group_ptrs(num_groups); + block_scale_ws.norm_constant = device_context.allocate_and_initialize_tensor( + options, + "norm_constant", + operation_desc.gemm.element_epilogue, + operation_desc.gemm.A.layout, // copied, but should this be D layout? + {1, 1}, + {1}, + 1, + seed_shift++, + 0 // device_index + ); + gemm_workspace_.block_scales.value().SFA_ptr_array_device.resize(gemm_workspace_.problem_count); + gemm_workspace_.block_scales.value().SFB_ptr_array_device.resize(gemm_workspace_.problem_count); + gemm_workspace_.block_scales.value().SFD_ptr_array_device.resize(gemm_workspace_.problem_count); for (size_t group_idx = 0; group_idx < num_groups; group_idx++) { - group_ptrs[group_idx] = input[group_idx]->batch_data(i); + auto group_str = std::to_string(group_idx); + block_scale_ws.SFA_ptr_array_host[group_idx] = device_context.create_ref_tensor( + options, + "block_SFA" + group_str, + block_scale_desc.SFA.element, + block_scale_desc.SFA.layout, + {sfa_m, sfa_sfb_k}, + {sfa_sfb_k}, + block_SFA->batch_data(group_idx), + 1, + 0); + block_scale_ws.SFB_ptr_array_host[group_idx] = device_context.create_ref_tensor( + options, + "block_SFB" + group_str, + block_scale_desc.SFB.element, + block_scale_desc.SFB.layout, + {sfb_n, sfa_sfb_k}, + {sfa_sfb_k}, + block_SFB->batch_data(group_idx), + 1, + 0); + block_scale_ws.SFD_ptr_array_host[group_idx] = device_context.create_ref_tensor( + options, + "block_SFD" + group_str, + block_scale_desc.SFD.element, + block_scale_desc.SFD.layout, + {sfd_m, sfd_n}, + {sfd_n}, + block_SFD->batch_data(group_idx), + 1, + 0); + block_scale_ws.SFD_reference_ptr_array_host[group_idx] = device_context.create_ref_tensor( + options, + "block_Reference_SFD" + group_str, + block_scale_desc.SFD.element, + block_scale_desc.SFD.layout, + {sfd_m, sfd_n}, + {sfd_n}, + block_ref_SFD->batch_data(group_idx), + 1, + 0); + } + for(int problem_idx = 0; problem_idx < gemm_workspace_.problem_count; problem_idx++) { + auto problem_str = std::to_string(problem_idx); + gemm_workspace_.block_scales.value().SFA_ptr_array_device[problem_idx] = device_context.create_ref_tensor( + options, + "block_SFA" + problem_str, + block_scale_desc.SFA.element, + block_scale_desc.SFA.layout, + {sfa_m, sfa_sfb_k}, + {sfa_sfb_k}, + block_SFA->batch_data(problem_idx*num_groups), + num_groups, + 0); + gemm_workspace_.block_scales.value().SFB_ptr_array_device[problem_idx] = device_context.create_ref_tensor( + options, + "block_SFB" + problem_str, + block_scale_desc.SFB.element, + block_scale_desc.SFB.layout, + {sfb_n, sfa_sfb_k}, + {sfa_sfb_k}, + block_SFB->batch_data(problem_idx*num_groups), + num_groups, + 0); + gemm_workspace_.block_scales.value().SFD_ptr_array_device[problem_idx] = device_context.create_ref_tensor( + options, + "block_SFD" + problem_str, + block_scale_desc.SFD.element, + block_scale_desc.SFD.layout, + {sfd_m, sfd_n}, + {sfd_n}, + block_SFD->batch_data(problem_idx*num_groups), + num_groups, + 0); } - dev_ptr_arrays[i]->copy_from_host(group_ptrs.data()); } - }; - create_dev_ptr_array_all_workspace( - gemm_workspace_.A_ptr_array_device, - gemm_workspace_.A_ptr_array_host, - "A"); - create_dev_ptr_array_all_workspace( - gemm_workspace_.B_ptr_array_device, - gemm_workspace_.B_ptr_array_host, - "B"); - create_dev_ptr_array_all_workspace( - gemm_workspace_.C_ptr_array_device, - gemm_workspace_.C_ptr_array_host, - "C"); - create_dev_ptr_array_all_workspace( - gemm_workspace_.D_ptr_array_device, - gemm_workspace_.D_ptr_array_host, - "D"); - if (is_block_scaled) { - auto& block_scale_ws = gemm_workspace_.block_scales.value(); - create_dev_ptr_array_all_workspace( - block_scale_ws.SFA_ptr_array_device, - block_scale_ws.SFA_ptr_array_host, - "SFA"); - create_dev_ptr_array_all_workspace( - block_scale_ws.SFB_ptr_array_device, - block_scale_ws.SFB_ptr_array_host, - "SFB"); - create_dev_ptr_array_all_workspace( - block_scale_ws.SFD_ptr_array_device, - block_scale_ws.SFD_ptr_array_host, - "SFD"); + for (size_t group_idx = 0; group_idx < num_groups; group_idx++) { + gemm_workspace_.tokens_per_expert_host[group_idx] = problem_.n(group_idx); + auto group_str = std::to_string(group_idx); - block_scale_ws.norm_constant = device_context.allocate_and_initialize_tensor( - options, - "norm_constant", - operation_desc.gemm.element_epilogue, - operation_desc.gemm.A.layout, // copied, but should this be D layout? - {1, 1}, - {1}, - 1, - seed_shift++, - 0 // device_index - ); - } - else if (is_blockwise) { - auto& block_scale_ws = gemm_workspace_.block_scales.value(); - create_dev_ptr_array_all_workspace( - block_scale_ws.SFA_ptr_array_device, - block_scale_ws.SFA_ptr_array_host, - "SFA"); - create_dev_ptr_array_all_workspace( - block_scale_ws.SFB_ptr_array_device, - block_scale_ws.SFB_ptr_array_host, - "SFB"); + gemm_workspace_.A_ptr_array_host[group_idx] = device_context.create_ref_tensor( + options, + "block_A" + group_str, + operation_desc.gemm.A.element, + operation_desc.gemm.A.layout, + {max_m, max_k}, + {int(problem_.max_lda)}, + block_A->batch_data(group_idx), + 1, + 0); + gemm_workspace_.B_ptr_array_host[group_idx] = device_context.create_ref_tensor( + options, + "block_B" + group_str, + operation_desc.gemm.B.element, + operation_desc.gemm.B.layout, + {max_k, max_n}, + {int(problem_.max_ldb)}, + block_B->batch_data(group_idx), + 1, + 0); + gemm_workspace_.C_ptr_array_host[group_idx] = device_context.create_ref_tensor( + options, + "block_C" + group_str, + operation_desc.gemm.C.element, + operation_desc.gemm.C.layout, + {max_m, max_n}, + {int(problem_.max_ldc)}, + block_C->batch_data(group_idx), + 1, + 0); + gemm_workspace_.D_ptr_array_host[group_idx] = device_context.create_ref_tensor( + options, + "block_D" + group_str, + operation_desc.gemm.D.element, + operation_desc.gemm.D.layout, + {max_m, max_n}, + {int(problem_.max_ldc)}, + block_D->batch_data(group_idx), + 1, + 0); + gemm_workspace_.reference_ptr_array_host[group_idx] = device_context.create_ref_tensor( + options, + "Reference_" + group_str, + operation_desc.gemm.D.element, + operation_desc.gemm.D.layout, + {max_m, max_n}, + {int(problem_.max_ldc)}, + block_ref_D->batch_data(group_idx), + 1, + 0); + + gemm_workspace_.A_ptr_array_device.resize(gemm_workspace_.problem_count); + gemm_workspace_.B_ptr_array_device.resize(gemm_workspace_.problem_count); + gemm_workspace_.C_ptr_array_device.resize(gemm_workspace_.problem_count); + gemm_workspace_.D_ptr_array_device.resize(gemm_workspace_.problem_count); + + for(int problem_idx = 0; problem_idx < gemm_workspace_.problem_count; problem_idx++) { + auto problem_str = std::to_string(problem_idx); + gemm_workspace_.A_ptr_array_device[problem_idx] = device_context.create_ref_tensor( + options, + "block_A" + problem_str, + operation_desc.gemm.A.element, + operation_desc.gemm.A.layout, + {max_m, max_k}, + {int(problem_.max_lda)}, + block_A->batch_data(problem_idx*num_groups), + num_groups, + 0); + gemm_workspace_.B_ptr_array_device[problem_idx] = device_context.create_ref_tensor( + options, + "block_B" + problem_str, + operation_desc.gemm.B.element, + operation_desc.gemm.B.layout, + {max_k, max_n}, + {int(problem_.max_ldb)}, + block_B->batch_data(problem_idx*num_groups), + num_groups, + 0); + gemm_workspace_.C_ptr_array_device[problem_idx] = device_context.create_ref_tensor( + options, + "block_C" + problem_str, + operation_desc.gemm.C.element, + operation_desc.gemm.C.layout, + {max_m, max_n}, + {int(problem_.max_ldc)}, + block_C->batch_data(problem_idx*num_groups), + num_groups, + 0); + gemm_workspace_.D_ptr_array_device[problem_idx] = device_context.create_ref_tensor( + options, + "block_D" + problem_str, + operation_desc.gemm.D.element, + operation_desc.gemm.D.layout, + {max_m, max_n}, + {int(problem_.max_ldc)}, + block_D->batch_data(problem_idx*num_groups), + num_groups, + 0); + + } + } + gemm_workspace_.tokens_per_expert_device->copy_from_host(gemm_workspace_.tokens_per_expert_host.data()); } init_arguments(options); diff --git a/tools/util/include/cutlass/util/reference/device/gemm_complex.h b/tools/util/include/cutlass/util/reference/device/gemm_complex.h index bddf5962..94cc5649 100644 --- a/tools/util/include/cutlass/util/reference/device/gemm_complex.h +++ b/tools/util/include/cutlass/util/reference/device/gemm_complex.h @@ -108,7 +108,9 @@ __global__ void GemmComplex( tensor_a.add_pointer_offset(batch_idx * batch_stride_A); tensor_b.add_pointer_offset(batch_idx * batch_stride_B); - tensor_c.add_pointer_offset(batch_idx * batch_stride_C); + if(tensor_c.data()) { + tensor_c.add_pointer_offset(batch_idx * batch_stride_C); + } tensor_d.add_pointer_offset(batch_idx * batch_stride_D); for (; batch_idx < batch_count; batch_idx += gridDim.z) { @@ -163,17 +165,20 @@ __global__ void GemmComplex( MatrixCoord coord = MatrixCoord(row, col); if (row < M && col < N) { - - tensor_d.at(coord) = convert_op( - alpha * ScalarType(accum[i][j]) + - beta * ScalarType(tensor_c.at(coord))); + ScalarType epilog = alpha * ScalarType(accum[i][j]); + if(tensor_c.data()) { + epilog += beta * ScalarType(tensor_c.at(coord)); + } + tensor_d.at(coord) = convert_op(epilog); } } } tensor_a.add_pointer_offset(batch_stride_A * gridDim.z); tensor_b.add_pointer_offset(batch_stride_B * gridDim.z); - tensor_c.add_pointer_offset(batch_stride_C * gridDim.z); + if(tensor_c.data()) { + tensor_c.add_pointer_offset(batch_stride_C * gridDim.z); + } tensor_d.add_pointer_offset(batch_stride_D * gridDim.z); } // for (batch_idx) diff --git a/tools/util/include/cutlass/util/reference/host/gemm_complex.h b/tools/util/include/cutlass/util/reference/host/gemm_complex.h index 221a6040..7ce1859f 100644 --- a/tools/util/include/cutlass/util/reference/host/gemm_complex.h +++ b/tools/util/include/cutlass/util/reference/host/gemm_complex.h @@ -154,10 +154,11 @@ void GemmComplex( MatrixCoord coord = MatrixCoord(row, col); if (row < M && col < N) { - - tensor_d.at(coord) = convert_op( - alpha * ScalarType(accum[i][j]) + - beta * ScalarType(tensor_c.at(coord))); + ScalarType epilog = alpha * ScalarType(accum[i][j]); + if(tensor_c.data()) { + epilog += beta * ScalarType(tensor_c.at(coord)); + } + tensor_d.at(coord) = convert_op(epilog); } } } @@ -167,7 +168,9 @@ void GemmComplex( tensor_a.add_pointer_offset(batch_stride_A); tensor_b.add_pointer_offset(batch_stride_B); - tensor_c.add_pointer_offset(batch_stride_C); + if(tensor_c.data()) { + tensor_c.add_pointer_offset(batch_stride_C); + } tensor_d.add_pointer_offset(batch_stride_D); } // for (batch_idx)