leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 24 Packages Projects Releases Wiki Activity

v4.2 tag release. (#2638)

Browse Source
This commit is contained in:
Junkai-Wu
2025-09-15 12:21:53 -04:00
committed by GitHub
parent 56f0718a97
commit 6a35b4d22f
161 changed files with 14056 additions and 3793 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+45 -3
View File
@@ -2,19 +2,46 @@
# CUTLASS 4.x
## [4.2.0](https://github.com/NVIDIA/cutlass/tree/main) (2025-08-21)
## [4.2.0](https://github.com/NVIDIA/cutlass/releases/tag/v4.2.0) (2025-09-15)
### CuTe DSL
* We will likely be skipping 4.2.dev release and directly target 4.2.
* CuTeDSL version remains at 4.1.0 till then.
* More Python versions are now supported for both x86-64 and aarch64, including
- Python 3.10, 3.11, 3.12, and 3.13
* Added new example and updated notebook to get started with CuTe DSL
- [Call kernels with dlpack bypassed](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/ampere/call_bypass_dlpack.py)
- Updates on [TensorSSA demonstration](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/tensorssa.ipynb)
+ Added a section for introducing the broadcast
* API updates
- Please refer to [DSL API changelog](https://docs.nvidia.com/cutlass/media/docs/pythonDSL/cute_dsl_api/changelog.html) for details
* Bug fixings and improvements
- Fixed ``cute.print_tensor`` for coordinate tensor
- Fixed `cute.print` for tuple of layouts
- Fixed frozen object is not properly updated after fully assigned in dynamic control flow
- Fixed assign tuple/list element in a dynamic control flow may cause compilation failure
- Improved error message when CUDA context is not initialized
- Improved docstring of congruent and weakly_congruent
### CUTLASS C++
* Add K major scale factor support for Hopper SM90 blockwise kernels.
* Further enhance Blackwell SM100 Attention kernels in [example 77](https://github.com/NVIDIA/cutlass/tree/main/examples/77_blackwell_fmha/).
- Add fused reduction kernel support for cutlass MLA.
- Add softmax skip correction.
- Support for GQA in FMHA backward kernel.
- Fix an issue where `get_unmasked_trip_count` may return a negative value.
- Fix an issue where mbarriers are initialized with a zero arrival count.
- Fix a corner case issue where the sequence length of q is not a multiple of tile_q.
- Remove tma padding for forward kernel inputs.
* Add Blackwell SM120 blockwise gemm kernel example: [example 87](https://github.com/NVIDIA/cutlass/tree/main/87_blackwell_geforce_gemm_blockwise/).
* Add Blackwell SM100 kernel example of MoE gemm using TMA+CPASYNC to load input matrices: [example 92](https://github.com/NVIDIA/cutlass/tree/main/examples/92_blackwell_moe_gemm/).
* Support for Blackwell SM103 kernels for B300 GPUs.
- Collective mainloop codes: [Blockscaled datatypes with support for dense GEMM mainloop](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/gemm/collective/sm103_blockscaled_mma_warpspecialized.hpp)
- New [GEMM](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/gemm/dispatch_policy.hpp) and [epilogue](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/epilogue/dispatch_policy.hpp) dispatch policies for collectives, kernel layers, and builders.
- Kernel codes: [Blockscaled datatypes with support for dense GEMM kernel](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_tma_warpspecialized.hpp).
* Set of examples that demonstrate the usage of the 3.x API for targeting Blackwell SM103 architecture:
- [Blockscaled ultra fp4 dense GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/89_sm103_fp4_ultra_gemm/).
- [Blockscaled ultra fp4 dense grouped GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/90_sm103_fp4_ultra_grouped_gemm).
* Set of unit tests that demonstrate the usage of Blackwell SM103 blockscaled GEMM
- Unit test files with prefix name of `sm103_` under [GEMM device unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/).
* Support for Blackwell SM100 cpasync kernel.
- Collective mainloop codes: [cpasync mainloop](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/gemm/collective/sm100_mma_cpasync_warpspecialized.hpp).
- Kernel codes: [cpasync kernel](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/gemm/kernel/sm100_gemm_cpasync_warpspecialized.hpp).
@@ -41,9 +68,24 @@
* Add support for heuristics-based kernel filtering and autotuning using `nvidia-matmul-heuristics`.
- Details please refer to [heuristics doc](https://github.com/NVIDIA/cutlass/tree/main/media/docs/cpp/heuristics.md).
* Rename legacy Python API package from `cutlass` to `cutlass_cppgen`.
* Add Blackwell EVT support to legacy Python interface.
- Restructuring the C++ Blackwell SM100 Collective Epilogue Builder to work with the Python interface's `EpilogueDescriptors`.
- Added Blackwell SM100 EVT Emitter on the Python side and routed most emission through Hopper SM90 Emitter.
- Added some support for running SM100 kernels via the Python interface.
* Instantiating more Blackwell kernels in profiler.
- Blackwell SM100 and SM103 kernels support `CUTLASS_LIBRARY_INSTANTIATION_LEVEL` to instantiate all possible combinations.
- To use this feature, `CUTLASS_LIBRARY_KERNELS` must be non-empty. Profiler will combine `CUTLASS_LIBRARY_KERNELS` and `CUTLASS_LIBRARY_INSTANTIATION_LEVEL` to instantiate specific kernels.
- Details please check [Profiler Doc](https://github.com/NVIDIA/cutlass/tree/main/media/docs/cpp/profiler.md).
* Fix some profiler issues:
- Modify default cluster callback values to none 0 to avoid profiler failure when these values are not set in command line.
- Fix some no output and timeout issues.
- Fix Pingpong Blockwise Hopper library generation.
* Fix some kernel issues:
- Fix Hopper SM90 group gemm kernel to only use the commit group and wait group instead of also waiting on mbarriers.
- Support Blackwell SM120 mixed input blockscaled grouped GEMM.
- Fix a tiny bug when K is large for Blackwell SM103 fp4 grouped GEMM kernel.
- Fix an issue in [example 68](https://github.com/NVIDIA/cutlass/tree/main/examples/68_hopper_fp8_warp_specialized_grouped_gemm_with_blockwise_scaling/) where problem size has value of 0.
- Relax k dimension constraints for blockwise gemm on Hopper in [example 68](https://github.com/NVIDIA/cutlass/tree/main/examples/68_hopper_fp8_warp_specialized_grouped_gemm_with_blockwise_scaling/).
* Add following unit tests:
- [fp16 accmulator for sm89 fp8 mma](https://github.com/NVIDIA/cutlass/tree/main/test/unit/cute/ampere/cooperative_gemm.cu)
- [movmatrix test](https://github.com/NVIDIA/cutlass/tree/main/test/unit/cute/turing/movm.cu)
CMakeLists.txt
+1 -1
View File
@@ -337,7 +337,7 @@ set(CUTLASS_LIBRARY_OPERATIONS "all" CACHE STRING "Comma-delimited list of opera
set(CUTLASS_LIBRARY_KERNELS ${CUTLASS_LIBRARY_KERNELS_INIT} CACHE STRING "Comma-delimited list of kernel name filters. If unspecified, only the largest tile size is enabled. If the string 'all' is specified, all kernels are enabled.")
set(CUTLASS_LIBRARY_IGNORE_KERNELS "" CACHE STRING "Comma-delimited list of kernels to exclude from build. This option ONLY takes effect if CUTLASS_LIBRARY_KERNELS is set.")
set(CUTLASS_LIBRARY_EXCLUDE_KERNELS "" CACHE STRING "Comma-delimited list of kernels to exclude from build. This option always takes effect, whether or not CUTLASS_LIBRARY_KERNELS is set. It also can exclude kernels from the filter file (see KERNEL_FILTER_FILE).")
set(CUTLASS_LIBRARY_INSTANTIATION_LEVEL "" CACHE STRING "Instantiation level for SM90 kernels. Set to `max` and make sure CUTLASS_LIBRARY_KERNELS is non-empty to stamp all possible kernel configurations.")
set(CUTLASS_LIBRARY_INSTANTIATION_LEVEL "" CACHE STRING "Instantiation level for SM90 and SM100 kernels. Set to `max` and make sure CUTLASS_LIBRARY_KERNELS is non-empty to stamp all possible kernel configurations.")
if(CUTLASS_LIBRARY_INSTANTIATION_LEVEL OR CUTLASS_LIBRARY_HEURISTICS_PROBLEMS_FILE)
message(STATUS "Enable extended SM90 WGMMA instruction shapes for instantiation levels")
README.md
+45 -3
View File
@@ -3,7 +3,7 @@
# CUTLASS 4.2.0
_CUTLASS 4.2.0 - Aug 2025_
_CUTLASS 4.2.0 - Sept 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,16 +46,43 @@ To get started quickly - please refer :
# What's New in CUTLASS 4.2
## CuTe DSL
* We will likely be skipping 4.2.dev release and directly target 4.2.
* CuTeDSL version remains at 4.1.0 till then.
* More Python versions are now supported for both x86-64 and aarch64, including
- Python 3.10, 3.11, 3.12, and 3.13
* Added new example and updated notebook to get started with CuTe DSL
- [Call kernels with dlpack bypassed](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/ampere/call_bypass_dlpack.py)
- Updates on [TensorSSA demonstration](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/tensorssa.ipynb)
+ Added a section for introducing the broadcast
* API updates
- Please refer to [DSL API changelog](https://docs.nvidia.com/cutlass/media/docs/pythonDSL/cute_dsl_api/changelog.html) for details
* Bug fixings and improvements
- Fixed ``cute.print_tensor`` for coordinate tensor
- Fixed `cute.print` for tuple of layouts
- Fixed frozen object is not properly updated after fully assigned in dynamic control flow
- Fixed assign tuple/list element in a dynamic control flow may cause compilation failure
- Improved error message when CUDA context is not initialized
- Improved docstring of congruent and weakly_congruent
## CUTLASS C++
* Add K major scale factor support for Hopper SM90 blockwise kernels.
* Further enhance Blackwell SM100 Attention kernels in [example 77](https://github.com/NVIDIA/cutlass/tree/main/examples/77_blackwell_fmha/).
- Add fused reduction kernel support for cutlass MLA.
- Add softmax skip correction.
- Support for GQA in FMHA backward kernel.
- Fix an issue where `get_unmasked_trip_count` may return a negative value.
- Fix an issue where mbarriers are initialized with a zero arrival count.
- Fix a corner case issue where the sequence length of q is not a multiple of tile_q.
- Remove tma padding for forward kernel inputs.
* Add Blackwell SM120 blockwise gemm kernel example: [example 87](https://github.com/NVIDIA/cutlass/tree/main/87_blackwell_geforce_gemm_blockwise/).
* Add Blackwell SM100 kernel example of MoE gemm using TMA+CPASYNC to load input matrices: [example 92](https://github.com/NVIDIA/cutlass/tree/main/examples/92_blackwell_moe_gemm/).
* Support for Blackwell SM103 kernels for B300 GPUs.
- Collective mainloop codes: [Blockscaled datatypes with support for dense GEMM mainloop](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/gemm/collective/sm103_blockscaled_mma_warpspecialized.hpp)
- New [GEMM](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/gemm/dispatch_policy.hpp) and [epilogue](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/epilogue/dispatch_policy.hpp) dispatch policies for collectives, kernel layers, and builders.
- Kernel codes: [Blockscaled datatypes with support for dense GEMM kernel](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_tma_warpspecialized.hpp).
* Set of examples that demonstrate the usage of the 3.x API for targeting Blackwell SM103 architecture:
- [Blockscaled ultra fp4 dense GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/89_sm103_fp4_ultra_gemm/).
- [Blockscaled ultra fp4 dense grouped GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/90_sm103_fp4_ultra_grouped_gemm).
* Set of unit tests that demonstrate the usage of Blackwell SM103 blockscaled GEMM
- Unit test files with prefix name of `sm103_` under [GEMM device unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/).
* Support for Blackwell SM100 cpasync kernel.
- Collective mainloop codes: [cpasync mainloop](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/gemm/collective/sm100_mma_cpasync_warpspecialized.hpp).
- Kernel codes: [cpasync kernel](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/gemm/kernel/sm100_gemm_cpasync_warpspecialized.hpp).
@@ -82,9 +109,24 @@ To get started quickly - please refer :
* Add support for heuristics-based kernel filtering and autotuning using `nvidia-matmul-heuristics`.
- Details please refer to [heuristics doc](https://github.com/NVIDIA/cutlass/tree/main/media/docs/cpp/heuristics.md).
* Rename legacy Python API package from `cutlass` to `cutlass_cppgen`.
* Add Blackwell EVT support to legacy Python interface.
- Restructuring the C++ Blackwell SM100 Collective Epilogue Builder to work with the Python interface's `EpilogueDescriptors`.
- Added Blackwell SM100 EVT Emitter on the Python side and routed most emission through Hopper SM90 Emitter.
- Added some support for running SM100 kernels via the Python interface.
* Instantiating more Blackwell kernels in profiler.
- Blackwell SM100 and SM103 kernels support `CUTLASS_LIBRARY_INSTANTIATION_LEVEL` to instantiate all possible combinations.
- To use this feature, `CUTLASS_LIBRARY_KERNELS` must be non-empty. Profiler will combine `CUTLASS_LIBRARY_KERNELS` and `CUTLASS_LIBRARY_INSTANTIATION_LEVEL` to instantiate specific kernels.
- Details please check [Profiler Doc](https://github.com/NVIDIA/cutlass/tree/main/media/docs/cpp/profiler.md).
* Fix some profiler issues:
- Modify default cluster callback values to none 0 to avoid profiler failure when these values are not set in command line.
- Fix some no output and timeout issues.
- Fix Pingpong Blockwise Hopper library generation.
* Fix some kernel issues:
- Fix Hopper SM90 group gemm kernel to only use the commit group and wait group instead of also waiting on mbarriers.
- Support Blackwell SM120 mixed input blockscaled grouped GEMM.
- Fix a tiny bug when K is large for Blackwell SM103 fp4 grouped GEMM kernel.
- Fix an issue in [example 68](https://github.com/NVIDIA/cutlass/tree/main/examples/68_hopper_fp8_warp_specialized_grouped_gemm_with_blockwise_scaling/) where problem size has value of 0.
- Relax k dimension constraints for blockwise gemm on Hopper in [example 68](https://github.com/NVIDIA/cutlass/tree/main/examples/68_hopper_fp8_warp_specialized_grouped_gemm_with_blockwise_scaling/).
* Add following unit tests:
- [fp16 accmulator for sm89 fp8 mma](https://github.com/NVIDIA/cutlass/tree/main/test/unit/cute/ampere/cooperative_gemm.cu)
- [movmatrix test](https://github.com/NVIDIA/cutlass/tree/main/test/unit/cute/turing/movm.cu)
examples/65_distributed_gemm/CMakeLists.txt
+2
View File
@@ -26,7 +26,9 @@
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
if (CUTLASS_NVCC_ARCHS MATCHES 90a)
cutlass_example_add_executable(
65_distributed_gemm
65_distributed_gemm.cu
)
endif()
examples/67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling/67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling.cu
+1 -1
View File
@@ -129,7 +129,7 @@ using ScaleConfig = decltype(cutlass::detail::sm90_trivial_blockwise_scale_confi
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::KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum;
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8Blockwise;
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
examples/67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling/67_hopper_fp8_warp_specialized_gemm_with_groupwise_scaling.cu
+1 -1
View File
@@ -137,7 +137,7 @@ using ScaleConfig = cutlass::detail::Sm90BlockwiseScaleConfig<ScaleGranularity
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::KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum;
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8Blockwise;
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltActAmaxAux<
examples/68_hopper_fp8_warp_specialized_grouped_gemm_with_blockwise_scaling/68_hopper_fp8_warp_specialized_grouped_gemm_with_blockwise_scaling.cu
+53 -34
View File
@@ -145,7 +145,7 @@ using ScaleConfig = cutlass::detail::Sm90BlockwiseScaleConfig<ScaleGranularity
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::KernelPtrArrayTmaWarpSpecializedCooperativeFP8BlockScaledAccum;
using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8Blockwise;
using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using FusionOperation = cutlass::epilogue::fusion::LinearCombination<ElementC, ElementAccumulator>;
@@ -402,12 +402,37 @@ void initialize(const OptionType &options) {
beta_host.clear();
for (int i = 0; i < options.groups; i++) {
ptr_A_host.at(i) = block_A.get() + offset_A.at(i);
ptr_B_host.at(i) = block_B.get() + offset_B.at(i);
ptr_C_host.at(i) = block_C.get() + offset_C.at(i);
ptr_D_host.at(i) = block_D.get() + offset_D.at(i);
ptr_blockscale_A_host.at(i) = blockscale_block_A.get() + offset_blockscale_A.at(i);
ptr_blockscale_B_host.at(i) = blockscale_block_B.get() + offset_blockscale_B.at(i);
// If the current group's matrix has size 0, set the pointer to nullptr
if (i < options.groups - 1 && offset_A.at(i) == offset_A.at(i + 1)) {
ptr_A_host.at(i) = nullptr;
} else {
ptr_A_host.at(i) = block_A.get() + offset_A.at(i);
}
if (i < options.groups - 1 && offset_B.at(i) == offset_B.at(i + 1)) {
ptr_B_host.at(i) = nullptr;
} else {
ptr_B_host.at(i) = block_B.get() + offset_B.at(i);
}
if (i < options.groups - 1 && offset_C.at(i) == offset_C.at(i + 1)) {
ptr_C_host.at(i) = nullptr;
} else {
ptr_C_host.at(i) = block_C.get() + offset_C.at(i);
}
if (i < options.groups - 1 && offset_D.at(i) == offset_D.at(i + 1)) {
ptr_D_host.at(i) = nullptr;
} else {
ptr_D_host.at(i) = block_D.get() + offset_D.at(i);
}
if (i < options.groups - 1 && offset_blockscale_A.at(i) == offset_blockscale_A.at(i + 1)) {
ptr_blockscale_A_host.at(i) = nullptr;
} else {
ptr_blockscale_A_host.at(i) = blockscale_block_A.get() + offset_blockscale_A.at(i);
}
if (i < options.groups - 1 && offset_blockscale_B.at(i) == offset_blockscale_B.at(i + 1)) {
ptr_blockscale_B_host.at(i) = nullptr;
} else {
ptr_blockscale_B_host.at(i) = blockscale_block_B.get() + offset_blockscale_B.at(i);
}
alpha_host.push_back((options.alpha == FLT_MAX) ? static_cast<ElementAccumulator>((rand() % 5) + 1) : options.alpha);
beta_host.push_back((options.beta == FLT_MAX) ? static_cast<ElementAccumulator>(rand() % 5) : options.beta);
ptr_alpha_host.at(i) = block_alpha.get() + i;
@@ -546,10 +571,10 @@ bool verify(const OptionType &options) {
blockscale_block_B.copy_to_host(blockscale_block_B_host.data());
bool passed = true;
std::cout << " Running host reference kernel - may run for a while for large problems." << std::endl;
for (int group_idx = 0; group_idx < options.groups; group_idx++) {
// Group scaling tensors shapes based `ScaleGranularityM`, CTA Block (TileShape) and GEMM Problem shape
auto [m, n, k] = options.problem_sizes_host.at(group_idx);
auto gemm_problem_shape = cute::make_shape(m, n, k);
// Create instantiation for device reference gemm kernel
auto A = cute::make_tensor(block_A_host.data() + offset_A.at(group_idx),
@@ -598,11 +623,7 @@ bool verify(const OptionType &options) {
ElementAccumulator,
ElementCompute,
decltype(C),
decltype(D),
unused_t, // bias
unused_t, // Aux
unused_t, // valpha
unused_t // vbeta
decltype(D)
> epilogue_params;
epilogue_params.C = C;
@@ -639,6 +660,24 @@ int run(OptionType &options, bool host_problem_shapes_available = true)
allocate(options);
initialize(options);
std::cout << " Problem Sizes, Alpha, Beta " << std::endl;
for (int32_t i = 0; i < options.groups; ++i) {
std::cout << " " << options.problem_sizes_host.at(i);
std::cout << ", " << alpha_host.at(i) << ", " << beta_host.at(i) << std::endl;
}
std::cout << " Groups : " << options.groups << std::endl;
std::cout << " Tile shape (M, N, K): " << size<0>(TileShape{}) << ", " << size<1>(TileShape{}) << ", " << size<2>(TileShape{}) << std::endl;
std::cout << " ScaleGranularityM: " << ScaleGranularityM << " (ScaleMsPerTile: " << ScaleMsPerTile << ")" << std::endl;
std::cout << " ScaleGranularityN: " << ScaleGranularityN << " (ScaleNsPerTile: " << ScaleNsPerTile << ")" << std::endl;
std::string raster = "Heuristic";
if (options.raster_order == RasterOrderOptions::AlongN) {
raster = "Along N";
}
else if (options.raster_order == RasterOrderOptions::AlongM) {
raster = "Along M";
}
std::cout << " Rasterization: " << raster << " with a maximum CTA swizzle of " << options.swizzle << std::endl;
// Instantiate CUTLASS kernel depending on templates
Gemm gemm;
@@ -671,8 +710,7 @@ int run(OptionType &options, bool host_problem_shapes_available = true)
}
// Run profiling loop
if (options.iterations > 0)
{
if (options.iterations > 0) {
GpuTimer timer;
timer.start();
for (int iter = 0; iter < options.iterations; ++iter) {
@@ -686,25 +724,6 @@ int run(OptionType &options, bool host_problem_shapes_available = true)
result.avg_runtime_ms = double(elapsed_ms) / double(options.iterations);
result.gflops = options.gflops(result.avg_runtime_ms / 1000.0);
std::string raster = "Heuristic";
if (options.raster_order == RasterOrderOptions::AlongN) {
raster = "Along N";
}
else if (options.raster_order == RasterOrderOptions::AlongM) {
raster = "Along M";
}
std::cout << " Problem Sizes, Alpha, Beta " << std::endl;
for (int32_t i = 0; i < options.groups; ++i) {
std::cout << " " << options.problem_sizes_host.at(i);
std::cout << ", " << alpha_host.at(i) << ", " << beta_host.at(i) << std::endl;
}
std::cout << " Groups : " << options.groups << std::endl;
std::cout << " Tile shape (M, N, K): " << size<0>(TileShape{}) << ", " << size<1>(TileShape{}) << ", " << size<2>(TileShape{}) << std::endl;
std::cout << " ScaleGranularityM: " << ScaleGranularityM << " (ScaleMsPerTile: " << ScaleMsPerTile << ")" << std::endl;
std::cout << " ScaleGranularityN: " << ScaleGranularityN << " (ScaleNsPerTile: " << ScaleNsPerTile << ")" << 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);
examples/68_hopper_fp8_warp_specialized_grouped_gemm_with_blockwise_scaling/68_hopper_fp8_warp_specialized_grouped_gemm_with_blockwise_scaling_with_sparse_groups.cu
+55 -31
View File
@@ -132,8 +132,7 @@ using ElementCompute = float; // E
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>; // This one is just to make the compiler happy with verify()...
using TileShape = Shape<_128,_128,_128>; // Threadblock-level tile size
using ClusterShape = Shape<_1,_1,_1>; // Shape of the threadblocks in a cluster
static constexpr int ScaleGranularityM = 1;
@@ -148,7 +147,7 @@ using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); // Layout ty
using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); // Layout type for SFB matrix operand
using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpongFP8BlockScaledAccum;
using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpongFP8Blockwise;
using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using FusionOperation = cutlass::epilogue::fusion::LinearCombination<ElementC, ElementAccumulator>;
@@ -407,12 +406,37 @@ void initialize(const OptionType &options) {
beta_host.clear();
for (int i = 0; i < options.groups; i++) {
ptr_A_host.at(i) = block_A.get() + offset_A.at(i);
ptr_B_host.at(i) = block_B.get() + offset_B.at(i);
ptr_C_host.at(i) = block_C.get() + offset_C.at(i);
ptr_D_host.at(i) = block_D.get() + offset_D.at(i);
ptr_blockscale_A_host.at(i) = blockscale_block_A.get() + offset_blockscale_A.at(i);
ptr_blockscale_B_host.at(i) = blockscale_block_B.get() + offset_blockscale_B.at(i);
// If the current group's matrix has size 0, set the pointer to nullptr
if (i < options.groups - 1 && offset_A.at(i) == offset_A.at(i + 1)) {
ptr_A_host.at(i) = nullptr;
} else {
ptr_A_host.at(i) = block_A.get() + offset_A.at(i);
}
if (i < options.groups - 1 && offset_B.at(i) == offset_B.at(i + 1)) {
ptr_B_host.at(i) = nullptr;
} else {
ptr_B_host.at(i) = block_B.get() + offset_B.at(i);
}
if (i < options.groups - 1 && offset_C.at(i) == offset_C.at(i + 1)) {
ptr_C_host.at(i) = nullptr;
} else {
ptr_C_host.at(i) = block_C.get() + offset_C.at(i);
}
if (i < options.groups - 1 && offset_D.at(i) == offset_D.at(i + 1)) {
ptr_D_host.at(i) = nullptr;
} else {
ptr_D_host.at(i) = block_D.get() + offset_D.at(i);
}
if (i < options.groups - 1 && offset_blockscale_A.at(i) == offset_blockscale_A.at(i + 1)) {
ptr_blockscale_A_host.at(i) = nullptr;
} else {
ptr_blockscale_A_host.at(i) = blockscale_block_A.get() + offset_blockscale_A.at(i);
}
if (i < options.groups - 1 && offset_blockscale_B.at(i) == offset_blockscale_B.at(i + 1)) {
ptr_blockscale_B_host.at(i) = nullptr;
} else {
ptr_blockscale_B_host.at(i) = blockscale_block_B.get() + offset_blockscale_B.at(i);
}
alpha_host.push_back((options.alpha == FLT_MAX) ? static_cast<ElementAccumulator>((rand() % 5) + 1) : options.alpha);
beta_host.push_back((options.beta == FLT_MAX) ? static_cast<ElementAccumulator>(rand() % 5) : options.beta);
ptr_alpha_host.at(i) = block_alpha.get() + i;
@@ -551,10 +575,10 @@ bool verify(const OptionType &options) {
blockscale_block_B.copy_to_host(blockscale_block_B_host.data());
bool passed = true;
std::cout << " Running host reference kernel - may run for a while for large problems." << std::endl;
for (int group_idx = 0; group_idx < options.groups; group_idx++) {
// Group scaling tensors shapes based `ScaleGranularityM`, CTA Block (TileShape) and GEMM Problem shape
auto [m, n, k] = options.problem_sizes_after_alignment_host.at(group_idx);
auto gemm_problem_shape = cute::make_shape(m, n, k);
// Create instantiation for device reference gemm kernel
auto A = cute::make_tensor(block_A_host.data() + offset_A.at(group_idx),
@@ -637,10 +661,27 @@ bool verify(const OptionType &options) {
template <typename OptionType>
int run(OptionType &options, bool host_problem_shapes_available = true)
{
allocate(options);
initialize(options);
std::cout << " Problem Sizes, Alpha, Beta " << std::endl;
for (int32_t i = 0; i < options.groups; ++i) {
std::cout << " " << options.problem_sizes_host.at(i);
std::cout << ", " << alpha_host.at(i) << ", " << beta_host.at(i) << std::endl;
}
std::cout << " Groups : " << options.groups << std::endl;
std::cout << " Tile shape (M, N, K): " << size<0>(TileShape{}) << ", " << size<1>(TileShape{}) << ", " << size<2>(TileShape{}) << std::endl;
std::cout << " ScaleGranularityM: " << ScaleGranularityM << " (ScaleMsPerTile: " << ScaleMsPerTile << ")" << std::endl;
std::cout << " ScaleGranularityN: " << ScaleGranularityN << " (ScaleNsPerTile: " << ScaleNsPerTile << ")" << std::endl;
std::string raster = "Heuristic";
if (options.raster_order == RasterOrderOptions::AlongN) {
raster = "Along N";
}
else if (options.raster_order == RasterOrderOptions::AlongM) {
raster = "Along M";
}
std::cout << " Rasterization: " << raster << " with a maximum CTA swizzle of " << options.swizzle << std::endl;
// Instantiate CUTLASS kernel depending on templates
Gemm gemm;
@@ -695,27 +736,10 @@ int run(OptionType &options, bool host_problem_shapes_available = true)
ScaleMsPerTile,
ScaleNsPerTile>(result.avg_runtime_ms / 1000.0);
std::string raster = "Heuristic";
if (options.raster_order == RasterOrderOptions::AlongN) {
raster = "Along N";
}
else if (options.raster_order == RasterOrderOptions::AlongM) {
raster = "Along M";
}
std::cout << " Problem Sizes, Alpha, Beta " << std::endl;
for (int32_t i = 0; i < options.groups; ++i) {
std::cout << " " << options.problem_sizes_host.at(i);
std::cout << ", " << alpha_host.at(i) << ", " << beta_host.at(i) << std::endl;
}
std::cout << " Groups : " << options.groups << std::endl;
std::cout << " Tile shape (M, N, K): " << size<0>(TileShape{}) << ", " << size<1>(TileShape{}) << ", " << size<2>(TileShape{}) << std::endl;
std::cout << " ScaleGranularityM: " << ScaleGranularityM << " (ScaleMsPerTile: " << ScaleMsPerTile << ")" << std::endl;
std::cout << " ScaleGranularityN: " << ScaleGranularityN << " (ScaleNsPerTile: " << ScaleNsPerTile << ")" << 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;
std::cout << " GBPS: " << result.gbps << std::endl;
fflush(stdout);
}
return 0;
@@ -766,8 +790,8 @@ int main(int argc, char const **args) {
// Evaluate CUTLASS kernels
//
std::cout << "Running tests with host problem shapes:" << std::endl;
run(options, true);
std::cout << "Running tests without host problem shapes:" << std::endl;
run(options, false);
examples/68_hopper_fp8_warp_specialized_grouped_gemm_with_blockwise_scaling/CMakeLists.txt
+5
View File
@@ -44,6 +44,9 @@ set(TEST_FIXED_LARGE_GROUP --m=2048 --n=512 --k=512 --groups=512 --iterations=0)
set(TEST_SMALL --m=256 --n=128 --iterations=0) # Small problem sizes
set(TEST_SMALL_LARGE_GROUP --m=128 --n=128 --groups=500 --iterations=0) # Small problem sizes
set(TEST_K_16B_ALIGNED --m=256 --n=512 --k=960 --groups=10 --iterations=0)
set(TEST_K_16B_ALIGNED_LARGE_GROUP --m=256 --n=512 --k=960 --groups=512 --iterations=0)
cutlass_example_add_executable(
68_hopper_fp8_warp_specialized_grouped_gemm_with_blockwise_scaling
68_hopper_fp8_warp_specialized_grouped_gemm_with_blockwise_scaling.cu
@@ -58,6 +61,8 @@ cutlass_example_add_executable(
TEST_FIXED_LARGE_GROUP
TEST_SMALL
TEST_SMALL_LARGE_GROUP
TEST_K_16B_ALIGNED
TEST_K_16B_ALIGNED_LARGE_GROUP
)
# MSVC will fail to compile this example with the following error:
examples/68_hopper_fp8_warp_specialized_grouped_gemm_with_blockwise_scaling/hopper_fp8_commandline.hpp
+6 -6
View File
@@ -111,14 +111,14 @@ struct Options {
int m = cmd_line_m;
int n = cmd_line_n;
int k = cmd_line_k;
if (m < 1) {
m = m_alignment * ((rand() % (64 * alignment / m_alignment)) + 1);
if (m < 0) {
m = m_alignment * (rand() % (64 * alignment / m_alignment));
}
if (n < 1) {
n = n_alignment * ((rand() % (64 * alignment / n_alignment)) + 1);
if (n < 0) {
n = n_alignment * (rand() % (64 * alignment / n_alignment));
}
if (k < 1) {
k = k_alignment * ((rand() % (32 * alignment / k_alignment)) + 1);
if (k < 0) {
k = k_alignment * (rand() % (32 * alignment / k_alignment));
}
problem_sizes_after_alignment_host.push_back({m, n, k});
problem_sizes_host.push_back({m, n, k});
examples/77_blackwell_fmha/77_blackwell_fmha.cu
+9 -9
View File
@@ -419,16 +419,16 @@ struct FwdRunner {
using ElementAccumulatorPV = float;
using ElementOut = cutlass::half_t;
// Q K D (B H)
// Q K D ((H_R, H_K) B)
using ProblemShapeRegular = cute::tuple<int, int, int, cute::tuple<cute::tuple<int, int>, int>>;
using ProblemShapeVarlen = cute::tuple<VariableLength, VariableLength, int, cute::tuple<cute::tuple<int, int>, int>>;
using ProblemShapeType = std::conditional_t<kIsVarlen, ProblemShapeVarlen, ProblemShapeRegular>;
using StrideQ = cute::tuple<int, _1, cute::tuple<cute::tuple<int, int>, int>>; // Q D (H_G H_R B)
using StrideK = cute::tuple<int, _1, cute::tuple<cute::tuple<_0, int>, int>>; // K D (H_G H_R B)
using StrideQ = cute::tuple<int, _1, cute::tuple<cute::tuple<int, int>, int>>; // Q D ((H_R, H_K), B)
using StrideK = cute::tuple<int, _1, cute::tuple<cute::tuple<_0, int>, int>>; // K D ((H_R, H_K), B)
using StrideV = StrideK;
using StrideO = StrideQ;
using StrideLSE = cute::tuple<_1, cute::tuple<cute::tuple<int, int>, int>>; // Q (H_G H_R B)
using StrideLSE = cute::tuple<_1, cute::tuple<cute::tuple<int, int>, int>>; // Q ((H_R, H_K), B)
static constexpr bool kIsPersistent = find_option_t<Tag::kIsPersistent, true_type, KernelOptions...>::value;
using TileScheduler = std::conditional_t<kIsPersistent, cutlass::fmha::kernel::PersistentTileScheduler, cutlass::fmha::kernel::IndividualTileScheduler>;
@@ -611,8 +611,8 @@ struct FwdRunner {
ProblemShapeType problem_size_for_launch;
get<0>(problem_size_for_launch) = VariableLength{max_seqlen_q};
get<1>(problem_size_for_launch) = VariableLength{max_seqlen_kv};
get<0>(problem_size_for_launch) = VariableLength{max_seqlen_q, nullptr, total_seqlen_q};
get<1>(problem_size_for_launch) = VariableLength{max_seqlen_kv, nullptr, total_seqlen_kv};
get<2>(problem_size_for_launch) = get<2>(problem_size);
get<3>(problem_size_for_launch) = get<3>(problem_size);
@@ -669,9 +669,9 @@ struct FwdRunner {
}
auto buffer_init_fn = [&](auto& buffer) {
buffer.block_Q.reset(size(shape_QO), kIsVarlen ? D*SQ*H : 0);
buffer.block_K.reset(size(shape_KV), kIsVarlen ? D*SK*H_K : 0);
buffer.block_V.reset(size(shape_KV), kIsVarlen ? D*SK*H_K : 0);
buffer.block_Q.reset(size(shape_QO));
buffer.block_K.reset(size(shape_KV));
buffer.block_V.reset(size(shape_KV));
buffer.block_O.reset(size(shape_QO), kIsVarlen ? D*SQ*H : 0);
buffer.block_LSE.reset(size(shape_LSE));
buffer.block_ref_O.reset(size(shape_QO), kIsVarlen ? D*SQ*H : 0);
examples/77_blackwell_fmha/77_blackwell_mla_fwd.cu
+7 -6
View File
@@ -590,8 +590,8 @@ struct MlaFwdRunner {
ProblemShapeType problem_size_for_launch;
get<0>(problem_size_for_launch) = VariableLength{max_seqlen_q};
get<1>(problem_size_for_launch) = VariableLength{max_seqlen_kv};
get<0>(problem_size_for_launch) = VariableLength{max_seqlen_q, nullptr, total_seqlen_q};
get<1>(problem_size_for_launch) = VariableLength{max_seqlen_kv, nullptr, total_seqlen_kv};
get<2>(problem_size_for_launch) = get<2>(problem_size);
get<3>(problem_size_for_launch) = get<3>(problem_size);
@@ -651,9 +651,9 @@ struct MlaFwdRunner {
}
auto buffer_init_fn = [&](auto& buffer) {
buffer.block_Q.reset(size(shape_Q), kIsVarlen ? D_latent_rope*SQ*H : 0);
buffer.block_K.reset(size(shape_K), kIsVarlen ? D_latent_rope*SK*H_K : 0);
buffer.block_V.reset(size(shape_V), kIsVarlen ? D*SK*H_K : 0);
buffer.block_Q.reset(size(shape_Q));
buffer.block_K.reset(size(shape_K));
buffer.block_V.reset(size(shape_V));
buffer.block_O.reset(size(shape_O), kIsVarlen ? D*SQ*H : 0);
buffer.block_LSE.reset(size(shape_LSE));
buffer.block_ref_O.reset(size(shape_O), kIsVarlen ? D*SQ*H : 0);
@@ -849,7 +849,8 @@ struct MlaFwdRunner {
flops *= static_cast<double>(size<3,1>(problem_shape));
}
flops *= 2.0 * (std::is_same_v<ActiveMask, CausalMask<false>> ? 0.5 : 1.0);
flops *= 2.0 * (std::is_same_v<ActiveMask, CausalMask<false>> ||
std::is_same_v<ActiveMask, CausalMask<true>> ? 0.5 : 1.0);
flops *= static_cast<double>(size<3,0>(problem_shape));
double flops0 = flops * static_cast<double>(size<2, 0>(problem_shape) + size<2, 1>(problem_shape));
examples/77_blackwell_fmha/CMakeLists.txt
+11 -3
View File
@@ -65,6 +65,8 @@ set(TEST_VARLEN_17 --verify --varlen --mask=causal --causal-type=qbegin --d=128
set(TEST_VARLEN_18 --verify --varlen --mask=causal --causal-type=qend --d=128 --h=4 --h_k=4 --varlen-q=128 --varlen-k=128)
set(TEST_VARLEN_19 --verify --varlen --mask=causal --causal-type=qend --d=128 --h=4 --h_k=4 --varlen-q=17 --varlen-k=257)
set(TEST_VARLEN_20 --verify --varlen --mask=causal --causal-type=qend --d=128 --h=4 --h_k=4 --varlen-q=17 --varlen-k=25)
set(TEST_VARLEN_21 --verify --varlen --mask=causal --causal-type=qend --d=128 --h=4 --h_k=4 --varlen-q=1013 --varlen-k=1024)
set(TEST_VARLEN_22 --verify --varlen --mask=causal --causal-type=qend --d=128 --h=4 --h_k=4 --varlen-q=1024 --varlen-k=1035)
@@ -89,6 +91,8 @@ set(TEST_MLA_FWD_VARLEN_17 --verify --varlen --mask=causal --causal-type=qbegin
set(TEST_MLA_FWD_VARLEN_18 --verify --varlen --mask=causal --causal-type=qend --d=128 --h=4 --h_k=4 --varlen-q=128 --varlen-k=128)
set(TEST_MLA_FWD_VARLEN_19 --verify --varlen --mask=causal --causal-type=qend --d=128 --h=4 --h_k=4 --varlen-q=17 --varlen-k=257)
set(TEST_MLA_FWD_VARLEN_20 --verify --varlen --mask=causal --causal-type=qend --d=128 --h=4 --h_k=4 --varlen-q=17 --varlen-k=25)
set(TEST_MLA_FWD_VARLEN_21 --verify --varlen --mask=causal --causal-type=qend --d=128 --h=4 --h_k=4 --varlen-q=1013 --varlen-k=1024)
set(TEST_MLA_FWD_VARLEN_22 --verify --varlen --mask=causal --causal-type=qend --d=128 --h=4 --h_k=4 --varlen-q=1024 --varlen-k=1035)
set(TEST_GEN_BASIC --b=1 --h=4 --k=512 --d=128 --verify)
@@ -140,6 +144,8 @@ if(NOT WIN32 AND (NOT (CMAKE_CXX_COMPILER_ID MATCHES "Clang")) AND (CUTLASS_NVCC
TEST_VARLEN_18
TEST_VARLEN_19
TEST_VARLEN_20
TEST_VARLEN_21
TEST_VARLEN_22
)
target_include_directories(77_blackwell_fmha_${PREC} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR})
target_compile_definitions(77_blackwell_fmha_${PREC} PRIVATE ${PREC_MACRO})
@@ -163,7 +169,7 @@ if(NOT WIN32 AND (NOT (CMAKE_CXX_COMPILER_ID MATCHES "Clang")) AND (CUTLASS_NVCC
77_blackwell_mla.cu
TEST_COMMAND_OPTIONS
TEST_MLA_BASIC
TEST_MLA_SEP_REDUCTION
TEST_MLA_SEP_REDUCTION
TEST_MLA_FUSE_REDUCTION
)
target_include_directories(77_blackwell_mla_2sm_${PREC} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR})
@@ -175,8 +181,8 @@ if(NOT WIN32 AND (NOT (CMAKE_CXX_COMPILER_ID MATCHES "Clang")) AND (CUTLASS_NVCC
77_blackwell_mla.cu
TEST_COMMAND_OPTIONS
TEST_MLA_BASIC
TEST_MLA_SEP_REDUCTION
TEST_MLA_FUSE_REDUCTION
TEST_MLA_SEP_REDUCTION
TEST_MLA_FUSE_REDUCTION
)
target_include_directories(77_blackwell_mla_2sm_cpasync_${PREC} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR})
target_compile_definitions(77_blackwell_mla_2sm_cpasync_${PREC} PRIVATE ${PREC_MACRO} CPASYNC)
@@ -241,6 +247,8 @@ if(NOT WIN32 AND (NOT (CMAKE_CXX_COMPILER_ID MATCHES "Clang")) AND (CUTLASS_NVCC
TEST_MLA_FWD_VARLEN_18
TEST_MLA_FWD_VARLEN_19
TEST_MLA_FWD_VARLEN_20
TEST_MLA_FWD_VARLEN_21
TEST_MLA_FWD_VARLEN_22
)
target_include_directories(77_blackwell_mla_fwd_${PREC} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR})
target_compile_definitions(77_blackwell_mla_fwd_${PREC} PRIVATE ${PREC_MACRO})
examples/77_blackwell_fmha/README.md
+3 -1
View File
@@ -8,7 +8,7 @@ For generation usage, use an M-blocking (Num-Groups) of 128 (although the limit
Context loads are done via TMA, whereas generation usage utilized `cp.async` and is thus more amenable to complex load patterns.
For variable sequence lenght, the code requires a batch of valid (but never used) padding memory ahead of the first input batch. This is achieved with least overhead by leaving one batch free and then arranging QKV consecutively.
For variable sequence length, the code requires a batch of valid (but never used) padding memory ahead of the first output batch. No padding is needed for the input tensor, but it requires that the input tensor contain no NaN or Inf values. Note that users should set `total_length` to the `problem_shape`.
The approach of this implementation is to reuse the selection logic of the collective gemm builder and recombine the result into an FMHA kernel.
The kernel and collective layer are then formulated to be fmha-specific.
@@ -67,6 +67,8 @@ For detailed information on how to invoke them, check out either the tests in `C
to simplify the sample, clarified that `fmha_gen` sample only supports head
dim 128.
* 4.3.0: For variable sequence length, the code requires a batch of valid (but never used) padding memory ahead of the first output batch. No padding is needed for the input tensor, but it requires that the input tensor contain no NaN or Inf values. Note that users should set `total_length` to the `problem_shape`.
# Copyright
Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
examples/77_blackwell_fmha/collective/fmha_fusion.hpp
+2 -2
View File
@@ -225,8 +225,8 @@ struct CausalMask : NoMask {
if constexpr (IsQBegin) {
return std::min(trip_count, int(ceil_div(size<0>(tile_shape), size<1>(tile_shape))));
} else {
const int offset_tile_q = get<1>(problem_size) % get<1>(tile_shape);
return std::min(trip_count, int(ceil_div(get<0>(tile_shape) + offset_tile_q, get<1>(tile_shape))));
const int corner_count = int((get<1>(problem_size) % get<1>(tile_shape) || get<0>(problem_size) % get<0>(tile_shape))) ;
return std::min(trip_count, int(ceil_div(get<0>(tile_shape), get<1>(tile_shape))) + corner_count);
}
}
examples/77_blackwell_fmha/collective/sm100_fmha_gen_mainloop_warpspecialized.hpp
+3 -3
View File
@@ -534,14 +534,14 @@ struct Sm100FmhaGenMainloopWarpspecialized {
tStS_v.data() = uint32_t(stage == _0{} ? TmemAllocation::V0 : TmemAllocation::V1);
Tensor tScS_v = tScS.compose(make_layout(make_shape(_128{}, _2{})));
auto tilePlikeFP32 = get<1>(TileShapeQK{}) / Int<sizeof(float)>{} * Int<sizeof(Element)>{};
auto tilePlikeFP32 = size<1>(TileShapeQK{}) / Int<sizeof(float)>{} * Int<sizeof(Element)>{};
Tensor tStS_P = tStS.compose(make_layout(make_shape(_128{}, tilePlikeFP32)));
tStS_P.data() = warp_uniform(uint32_t(stage == _0{} ? TmemAllocation::P0 : TmemAllocation::P1));
Tensor tScS_P = tScS.compose(make_layout(make_shape(_128{}, tilePlikeFP32)));
// Each thread owns a single row
using TMEM_LOAD = SM100_TMEM_LOAD_32dp32b32x; // 4x32 threads with 128 cols of 32b elem
using TMEM_STORE = SM100_TMEM_STORE_32dp32b32x; // 4x32 threads with 128 cols of 8b elem
using TMEM_LOAD = conditional_t<size<1>(TileShapeQK{}) < _128{}, SM100_TMEM_LOAD_32dp32b8x, SM100_TMEM_LOAD_32dp32b32x>; // 4x32 threads with 128 cols of 8b elem
using TMEM_STORE = conditional_t<size<1>(TileShapeQK{}) < _128{}, SM100_TMEM_STORE_32dp32b8x, SM100_TMEM_STORE_32dp32b32x>; // 4x32 threads with 128 cols of 8b elem
using TMEM_STORE_V = SM100_TMEM_STORE_32dp32b2x; // 4x32 threads with 2 cols of 32b elem
int thread_idx = threadIdx.x % (4 * cutlass::NumThreadsPerWarp);
examples/77_blackwell_fmha/collective/sm100_fmha_load_tma_warpspecialized.hpp
+16 -33
View File
@@ -95,32 +95,21 @@ struct Sm100FmhaLoadTmaWarpspecialized {
auto dQ = args.dQ;
auto dK = args.dK;
auto dV = args.dV;
auto problem_shape_qk = problem_shape;
using IntProblemShape = cute::tuple<int, int, int, cute::tuple<cute::tuple<int, int>, int>>;
IntProblemShape problem_shape_qk;
if constexpr (is_variable_length_v<tuple_element_t<0, ProblemShape>>) {
auto cumulative_length_q = get<0>(problem_shape).cumulative_length;
if (cumulative_length_q != nullptr) {
int max_length_q = get<0>(problem_shape).max_length;
// for variable sequence lenght, the batch is in units of row_stride
get<2,1>(dQ) = get<0>(dQ);
get<3,1>(problem_shape_qk) = std::max(get<3,1>(problem_shape_qk), max_length_q * (1 + get<3,1>(problem_shape)));
// offset ptr by the amount we add back in later
ptr_Q -= max_length_q * get<0>(dQ);
}
}
if constexpr (is_variable_length_v<tuple_element_t<1, ProblemShape>>) {
auto cumulative_length_kv = get<1>(problem_shape).cumulative_length;
if (cumulative_length_kv != nullptr) {
int max_length_kv = get<1>(problem_shape).max_length;
// for variable sequence lenght, the batch is in units of row_stride
get<2,1>(dK) = get<0>(dK);
get<2,1>(dV) = get<0>(dV);
get<3,1>(problem_shape_qk) = std::max(get<3,1>(problem_shape_qk), max_length_kv * (1 + get<3,1>(problem_shape)));
// offset ptr by the amount we add back in later
ptr_K -= max_length_kv * get<0>(dK);
ptr_V -= max_length_kv * get<0>(dV);
auto cumulative_length_k = get<1>(problem_shape).cumulative_length;
if (cumulative_length_q != nullptr && cumulative_length_k != nullptr ) {
get<0>(problem_shape_qk) = get<0>(problem_shape).total_length;
get<1>(problem_shape_qk) = get<1>(problem_shape).total_length;
get<2>(problem_shape_qk) = get<2>(problem_shape);
get<3>(problem_shape_qk) = get<3>(problem_shape);
}
} else {
problem_shape_qk = problem_shape;
}
auto params_qk = CollectiveMmaQK::to_underlying_arguments(
@@ -181,19 +170,16 @@ struct Sm100FmhaLoadTmaWarpspecialized {
Tensor mQ_qdl_p = params.tma_load_q.get_tma_tensor(select<0,2,3>(problem_shape));
int q_offs_0 = 0;
int q_offs_2_1 = 0;
if constexpr (is_variable_length_v<tuple_element_t<0, ParamsProblemShape>>) {
auto cumulative_length_q = get<0>(params_problem_shape).cumulative_length;
if (cumulative_length_q != nullptr) {
int max_length_q = get<0>(params_problem_shape).max_length;
q_offs_0 = max_length_q - get<0>(problem_shape);
q_offs_2_1 = cumulative_length_q[get<2,1>(blk_coord_q)] + get<0>(problem_shape);
q_offs_0 = cumulative_length_q[get<2,1>(blk_coord_q)];
get<2,1>(blk_coord_q) = 0;
}
}
Tensor mQ_qdl = domain_offset(make_coord(q_offs_0, _0{}, make_coord(_0{}, q_offs_2_1)), mQ_qdl_p);
Tensor mQ_qdl = domain_offset(make_coord(q_offs_0, _0{}, make_coord(_0{}, _0{})), mQ_qdl_p);
Tensor gQ_qdl = local_tile(mQ_qdl, TileShapeQK{}, make_coord(_, _, _), Step<_1, X, _1>{});
Tensor tSgQ_qdl = mma_qk.partition_A(gQ_qdl);
@@ -208,19 +194,16 @@ struct Sm100FmhaLoadTmaWarpspecialized {
Tensor mK_kdl_p = params.tma_load_k.get_tma_tensor(select<1,2,3>(problem_shape));
int kv_offs_0 = 0;
int kv_offs_2_1 = 0;
if constexpr (is_variable_length_v<tuple_element_t<1, ParamsProblemShape>>) {
auto cumulative_length = get<1>(params_problem_shape).cumulative_length;
if (cumulative_length != nullptr) {
int max_length = get<1>(params_problem_shape).max_length;
kv_offs_0 = max_length - get<1>(problem_shape);
kv_offs_2_1 = cumulative_length[get<2,1>(blk_coord_kv)] + get<1>(problem_shape);
kv_offs_0 = cumulative_length[get<2,1>(blk_coord_kv)];
get<2,1>(blk_coord_kv) = 0;
}
}
Tensor mK_kdl = domain_offset(make_coord(kv_offs_0, _0{}, make_coord(_0{}, kv_offs_2_1)), mK_kdl_p);
Tensor mK_kdl = domain_offset(make_coord(kv_offs_0, _0{}, make_coord(_0{}, _0{})), mK_kdl_p);
Tensor gK_kdl = local_tile(mK_kdl, TileShapeQK{}, make_coord(_, _, _), Step<X, _1, _1>{});
Tensor tSgK_kdl = mma_qk.partition_B(gK_kdl);
@@ -235,7 +218,7 @@ struct Sm100FmhaLoadTmaWarpspecialized {
ThrMMA mma_pv = typename CollectiveMmaPV::TiledMma{}.get_slice(0);
Tensor mV_dkl_p = params.tma_load_v.get_tma_tensor(select<2,1,3>(problem_shape));
Tensor mV_dkl = domain_offset(make_coord(_0{}, kv_offs_0, make_coord(_0{}, kv_offs_2_1)), mV_dkl_p);
Tensor mV_dkl = domain_offset(make_coord(_0{}, kv_offs_0, make_coord(_0{}, _0{})), mV_dkl_p);
Tensor gV_dkl = local_tile(mV_dkl, TileShapePV{}, make_coord(_, _, _), Step<X, _1, _1>{});
Tensor tOgV_dkl = mma_pv.partition_B(gV_dkl);
examples/77_blackwell_fmha/collective/sm100_fmha_mla_load_tma_warpspecialized.hpp
+16 -33
View File
@@ -102,32 +102,21 @@ struct Sm100MlaFwdLoadTmaWarpspecialized {
auto dQ = args.dQ;
auto dK = args.dK;
auto dV = args.dV;
auto problem_shape_qk = replace<2>(problem_shape, get<2, 0>(problem_shape) + get<2, 1>(problem_shape));
using IntProblemShape = cute::tuple<int, int, int, cute::tuple<cute::tuple<int, int>, int>>;
IntProblemShape problem_shape_qk;
if constexpr (is_variable_length_v<tuple_element_t<0, ProblemShape>>) {
auto cumulative_length_q = get<0>(problem_shape).cumulative_length;
if (cumulative_length_q != nullptr) {
int max_length_q = get<0>(problem_shape).max_length;
// for variable sequence lenght, the batch is in units of row_stride
get<2,1>(dQ) = get<0>(dQ);
get<3,1>(problem_shape_qk) = std::max(get<3,1>(problem_shape_qk), max_length_q * (1 + get<3,1>(problem_shape)));
// offset ptr by the amount we add back in later
ptr_Q -= max_length_q * get<0>(dQ);
}
}
if constexpr (is_variable_length_v<tuple_element_t<1, ProblemShape>>) {
auto cumulative_length_kv = get<1>(problem_shape).cumulative_length;
if (cumulative_length_kv != nullptr) {
int max_length_kv = get<1>(problem_shape).max_length;
// for variable sequence lenght, the batch is in units of row_stride
get<2,1>(dK) = get<0>(dK);
get<2,1>(dV) = get<0>(dV);
get<3,1>(problem_shape_qk) = std::max(get<3,1>(problem_shape_qk), max_length_kv * (1 + get<3,1>(problem_shape)));
// offset ptr by the amount we add back in later
ptr_K -= max_length_kv * get<0>(dK);
ptr_V -= max_length_kv * get<0>(dV);
auto cumulative_length_k = get<1>(problem_shape).cumulative_length;
if (cumulative_length_q != nullptr && cumulative_length_k != nullptr ) {
get<0>(problem_shape_qk) = get<0>(problem_shape).total_length;
get<1>(problem_shape_qk) = get<1>(problem_shape).total_length;
get<2>(problem_shape_qk) = get<2, 0>(problem_shape) + get<2, 1>(problem_shape);
get<3>(problem_shape_qk) = get<3>(problem_shape);
}
} else {
problem_shape_qk = replace<2>(problem_shape, get<2, 0>(problem_shape) + get<2, 1>(problem_shape));;
}
auto problem_shape_pv = replace<1>(select<0,2,1,3>(problem_shape_qk), get<2, 0>(problem_shape));
@@ -192,19 +181,16 @@ struct Sm100MlaFwdLoadTmaWarpspecialized {
Tensor mQ_qdl_p = params.tma_load_q.get_tma_tensor(select<0,2,3>(problem_shape_qk));
int q_offs_0 = 0;
int q_offs_2_1 = 0;
if constexpr (is_variable_length_v<tuple_element_t<0, ParamsProblemShape>>) {
auto cumulative_length_q = get<0>(params_problem_shape).cumulative_length;
if (cumulative_length_q != nullptr) {
int max_length_q = get<0>(params_problem_shape).max_length;
q_offs_0 = max_length_q - get<0>(problem_shape);
q_offs_2_1 = cumulative_length_q[get<2,1>(blk_coord_q)] + get<0>(problem_shape);
q_offs_0 = cumulative_length_q[get<2,1>(blk_coord_q)];
get<2,1>(blk_coord_q) = 0;
}
}
Tensor mQ_qdl = domain_offset(make_coord(q_offs_0, _0{}, make_coord(_0{}, q_offs_2_1)), mQ_qdl_p);
Tensor mQ_qdl = domain_offset(make_coord(q_offs_0, _0{}, make_coord(_0{}, _0{})), mQ_qdl_p);
Tensor gQ_qdl = local_tile(mQ_qdl, TileShapeQK{}, make_coord(_, _, _), Step<_1, X, _1>{});
Tensor tSgQ_qdl = mma_qk.partition_A(gQ_qdl);
@@ -219,19 +205,16 @@ struct Sm100MlaFwdLoadTmaWarpspecialized {
Tensor mK_kdl_p = params.tma_load_k.get_tma_tensor(select<1,2,3>(problem_shape_qk));
int kv_offs_0 = 0;
int kv_offs_2_1 = 0;
if constexpr (is_variable_length_v<tuple_element_t<1, ParamsProblemShape>>) {
auto cumulative_length = get<1>(params_problem_shape).cumulative_length;
if (cumulative_length != nullptr) {
int max_length = get<1>(params_problem_shape).max_length;
kv_offs_0 = max_length - get<1>(problem_shape);
kv_offs_2_1 = cumulative_length[get<2,1>(blk_coord_kv)] + get<1>(problem_shape);
kv_offs_0 = cumulative_length[get<2,1>(blk_coord_kv)];
get<2,1>(blk_coord_kv) = 0;
}
}
Tensor mK_kdl = domain_offset(make_coord(kv_offs_0, _0{}, make_coord(_0{}, kv_offs_2_1)), mK_kdl_p);
Tensor mK_kdl = domain_offset(make_coord(kv_offs_0, _0{}, make_coord(_0{}, _0{})), mK_kdl_p);
Tensor gK_kdl = local_tile(mK_kdl, TileShapeQK{}, make_coord(_, _, _), Step<X, _1, _1>{});
Tensor tSgK_kdl = mma_qk.partition_B(gK_kdl);
@@ -246,7 +229,7 @@ struct Sm100MlaFwdLoadTmaWarpspecialized {
ThrMMA mma_pv = typename CollectiveMmaPV::TiledMma{}.get_slice(0);
Tensor mV_dkl_p = params.tma_load_v.get_tma_tensor(select<2,1,3>(problem_shape_v));
Tensor mV_dkl = domain_offset(make_coord(_0{}, kv_offs_0, make_coord(_0{}, kv_offs_2_1)), mV_dkl_p);
Tensor mV_dkl = domain_offset(make_coord(_0{}, kv_offs_0, make_coord(_0{}, _0{})), mV_dkl_p);
Tensor gV_dkl = local_tile(mV_dkl, TileShapePV{}, make_coord(_, _, _), Step<X, _1, _1>{});
Tensor tOgV_dkl = mma_pv.partition_B(gV_dkl);
examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp
+5
View File
@@ -1215,6 +1215,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
Tensor tTR_cST_p = thread_t2r.partition_D(cST);
Tensor tTR_cST = split_wg(tTR_cST_p);
Tensor tTR_rST = make_tensor<ElementAcc>(shape(tTR_cST));
// Tensor tTR_tST_p = thread_t2r.partition_S(tSTtST);
Tensor tTR_tST = split_wg(thread_t2r.partition_S(tSTtST));
Tensor tTR_cDPT_p = thread_t2r.partition_D(cDPT);
@@ -1507,6 +1508,9 @@ 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");
#else
int warp_idx = cutlass::canonical_warp_idx_sync();
auto role = warp_idx_to_role(warp_idx);
uint32_t lane_predicate = cute::elect_one_sync();
@@ -1835,6 +1839,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
/* no-op */
}
#endif
}
static dim3 get_block_shape() {
examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp
+4
View File
@@ -1480,6 +1480,9 @@ 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");
#else
int warp_idx = cutlass::canonical_warp_idx_sync();
auto role = warp_idx_to_role(warp_idx);
uint32_t lane_predicate = cute::elect_one_sync();
@@ -1804,6 +1807,7 @@ struct Sm100FmhaBwdMlaKernelTmaWarpSpecialized {
/* no-op */
}
#endif
}
static dim3 get_block_shape() {
examples/77_blackwell_fmha/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp
+27 -6
View File
@@ -251,6 +251,9 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
}
CUTLASS_DEVICE void operator()(const Params &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");
#else
TileScheduler tile_scheduler{params.tile_scheduler};
@@ -465,6 +468,8 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
else if (role == WarpRole::Correction) {
cutlass::arch::warpgroup_reg_dealloc<NumRegsCorrection>();
bool has_valid = false;
CUTLASS_PRAGMA_NO_UNROLL
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
auto blk_coord = tile_scheduler.get_block_coord();
@@ -476,6 +481,8 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
continue;
}
has_valid = true;
if (get<1>(logical_problem_shape) == 0) {
mainloop.correction_empty(
blk_coord,
@@ -505,16 +512,17 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
if constexpr (NumWarpsEpilogue == 0) {
static_assert(NumWarpsCorrection == 1);
uint32_t free_stage_ptr = shared_storage.tmem_base_ptr;
tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns);
if (has_valid) {
uint32_t free_stage_ptr = shared_storage.tmem_base_ptr;
tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns);
}
}
}
else if (role == WarpRole::MMA) {
warpgroup_reg_set<NumRegsOther>();
tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);
__syncwarp();
bool allocated = false;
CUTLASS_PRAGMA_NO_UNROLL
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
@@ -527,6 +535,12 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
continue;
}
if (!allocated) {
tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);
__syncwarp();
allocated = true;
}
if (get<1>(logical_problem_shape) == 0) {
continue;
}
@@ -580,6 +594,8 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
else if (role == WarpRole::Epilogue) {
warpgroup_reg_set<NumRegsOther>();
bool has_valid = false;
CUTLASS_PRAGMA_NO_UNROLL
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
auto blk_coord = tile_scheduler.get_block_coord();
@@ -591,6 +607,8 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
continue;
}
has_valid = true;
epilogue.store(
blk_coord, logical_problem_shape,
params.epilogue, params.problem_shape,
@@ -602,8 +620,10 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
static_assert(NumWarpsEpilogue <= 1);
if constexpr (NumWarpsEpilogue == 1) {
uint32_t free_stage_ptr = shared_storage.tmem_base_ptr;
tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns);
if(has_valid) {
uint32_t free_stage_ptr = shared_storage.tmem_base_ptr;
tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns);
}
}
}
@@ -612,6 +632,7 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
/* no-op, donate regs and exit */
}
#endif
}
};
examples/77_blackwell_fmha/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp
+4
View File
@@ -247,6 +247,9 @@ struct Sm100FmhaGenKernelWarpspecialized {
}
CUTLASS_DEVICE void operator()(const Params &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");
#else
TileScheduler tile_scheduler{params.tile_scheduler};
@@ -569,6 +572,7 @@ struct Sm100FmhaGenKernelWarpspecialized {
/* no-op, donate regs and exit */
}
#endif
}
};
examples/77_blackwell_fmha/kernel/sm100_fmha_mla_tma_warpspecialized.hpp
+4
View File
@@ -507,6 +507,9 @@ 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");
#else
TileScheduler tile_scheduler(params.tile_scheduler);
@@ -814,6 +817,7 @@ struct Sm100FmhaMlaKernelTmaWarpspecialized {
uint32_t free_stage_ptr = shared_storage.tmem_base_ptr;
tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns);
}
#endif
}
template<class BlkCoord>
examples/82_blackwell_distributed_gemm/CMakeLists.txt
+2
View File
@@ -26,7 +26,9 @@
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
if (CUTLASS_NVCC_ARCHS MATCHES 100a)
cutlass_example_add_executable(
82_blackwell_distributed_gemm
82_blackwell_distributed_gemm.cu
)
endif()
examples/86_blackwell_mixed_dtype_gemm/86_blackwell_mixed_dtype.cu
+4 -2
View File
@@ -331,11 +331,13 @@ bool verify(MixedDtypeOptions const& options) {
//
// Compute reference output
//
constexpr int AlignmentBdq = 128 / cutlass::sizeof_bits<MmaType>::value;
using CollectiveMainloopRef = typename cutlass::gemm::collective::CollectiveBuilder<
ArchTag, OperatorClass,
MmaType, LayoutA, AlignmentA,
MmaType, LayoutB, AlignmentB,
MmaType, LayoutB, AlignmentBdq,
ElementAccumulator,
MmaTileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
examples/88_hopper_fmha/kernel/fmha_kernel_tma.hpp
+4
View File
@@ -137,6 +137,9 @@ struct FmhaKernelTma {
}
CUTLASS_DEVICE void operator()(const Params &params, char* smem) {
#if ! defined(CUTLASS_ARCH_MMA_SM90A_ENABLED)
printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
#else
TileScheduler tile_scheduler{params.tile_scheduler};
// Shared memory.
@@ -216,6 +219,7 @@ struct FmhaKernelTma {
result, typename CollectiveMainloop::TiledMmaPV{},
params.problem_size, params.epilogue,
epi_load_pipeline, storage.epilogue);
#endif
}
};
examples/88_hopper_fmha/kernel/fmha_kernel_tma_warpspecialized.hpp
+4
View File
@@ -160,6 +160,9 @@ struct FmhaKernelTmaWarpSpecialized {
}
CUTLASS_DEVICE void operator()(const Params &params, char* smem) {
#if ! defined(CUTLASS_ARCH_MMA_SM90A_ENABLED)
printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
#else
enum class WarpGroupRole {
Producer = 0,
@@ -412,6 +415,7 @@ struct FmhaKernelTmaWarpSpecialized {
if constexpr (kIsEpilogueLocked) ; math_wg_order_barrier.arrive();
}
}
#endif
}
};
examples/90_sm103_fp4_ultra_grouped_gemm/CMakeLists.txt
+24
View File
@@ -26,6 +26,21 @@
# 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_RANDOM --iterations=0) # Random problem sizes
set(TEST_EPILOGUE --alpha=0.5 --beta=0.5 --iterations=0) # Random problem sizes
set(TEST_EPILOGUE_LARGE_GROUP --alpha=1.5 --beta=2.0 --groups=50 --iterations=0) # Random problem sizes
set(TEST_EPILOGUE_OP --beta=0.5 --iterations=1) # Random problem sizes
set(TEST_FIXED --m=2048 --n=5120 --k=8192 --iterations=0) # Fixed problem sizes
set(TEST_FIXED_LARGE_GROUP --m=2048 --n=512 --beta=2.0 --k=512 --groups=51 --iterations=0)
set(TEST_SMALL --m=256 --n=128 --iterations=0) # Small problem sizes
set(TEST_SMALL_LARGE_GROUP --m=128 --n=128 --beta=0.5 --groups=50 --iterations=0) # Small problem sizes
set(TEST_RANDOM_PERF --iterations=10) # Random problem sizes
set(TEST_RANDOM_SMALL_GROUP --groups=3 --iterations=1) # Random problem sizes
set(TEST_EPILOGUE_SMALL_GROUP --alpha=1.5 --beta=2.0 --groups=3 --iterations=1) # Random problem sizes
@@ -35,6 +50,15 @@ cutlass_example_add_executable(
90_sm103_fp4_ultra_grouped_gemm
90_sm103_fp4_ultra_grouped_gemm.cu
TEST_COMMAND_OPTIONS
TEST_RANDOM
TEST_EPILOGUE
TEST_EPILOGUE_LARGE_GROUP
TEST_EPILOGUE_OP
TEST_FIXED
TEST_FIXED_LARGE_GROUP
TEST_SMALL
TEST_SMALL_LARGE_GROUP
TEST_RANDOM_PERF
TEST_RANDOM_SMALL_GROUP
TEST_EPILOGUE_SMALL_GROUP
)
examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_fp4_grouped.cu
+701
View File
@@ -0,0 +1,701 @@
/***************************************************************************************************
* Copyright (c) 2025 - 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.
*
**************************************************************************************************/
/*! \file
\brief Example of Blackwell MoE-style grouped NVFP4 GEMM implementation using TMA to load A and CPASYNC to load B.
This example demonstrates an implementation of GEMM using mixed TMA+CPASYNC to load input matrices.
In the decoding stage of Mixture of Experts (MoE) models, the number of tokens in different experts
can varies a lot, which requires frequently updates of TMA descriptors in TMA-based implementation.
This examples uses CPASYNC to load activation (B) matrix to avoid the overhead of updating TMA descriptors.
Usage:
$ ./examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_fp4_grouped
--m=28672 --n=4 --k=4096 --l=8 --benchmark=benchmark.txt
*/
#include <iostream>
#include <fstream>
#include "cute/tensor.hpp"
#include "cutlass/cutlass.h"
#include "cutlass/tensor_ref.h"
#include "cutlass/epilogue/collective/default_epilogue.hpp"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/gemm/kernel/tile_scheduler.hpp"
#include "cute/tensor.hpp"
#include "cutlass/tensor_ref.h"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/detail/sm100_blockscaled_layout.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/gemm/kernel/tile_scheduler_params.h"
#include "cutlass/util/command_line.h"
#include "cutlass/util/distribution.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/packed_stride.hpp"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/util/reference/device/gemm.h"
#include "cutlass/util/reference/device/tensor_compare.h"
#include "cutlass/util/reference/host/tensor_fill.h"
#include "cutlass/util/reference/host/gett.hpp"
#include "cutlass/util/reference/host/tensor_norm.h"
#include "cutlass/util/reference/host/tensor_compare.h"
#include "cutlass/util/command_line.h"
#include "cutlass/util/distribution.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/packed_stride.hpp"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/util/reference/device/gemm_complex.h"
#include "cutlass/util/reference/device/tensor_compare.h"
#include "cutlass/util/reference/device/tensor_fill.h"
#include "cutlass/util/reference/device/gemm.h"
#include "cutlass/util/reference/device/tensor_compare.h"
#include "cutlass/util/reference/host/tensor_fill.h"
#include "cutlass/util/reference/host/gett.hpp"
#include "cutlass/util/reference/host/tensor_norm.h"
#include "cutlass/util/reference/host/tensor_compare.h"
#include "helper.h"
using namespace cute;
///////////////////////////////////////////////////////////////////////////////////////////////////
/// Command line options parsing
struct Options {
bool help;
bool error;
bool verification;
int m, n, k, l;
int iterations;
std::string benchmark_path;
Options():
help(false),
error(false),
verification(true),
m(2048), n(2048), k(2048), l(1),
iterations(10)
{ }
// Parses the command line
void parse(int argc, char const **args) {
cutlass::CommandLine cmd(argc, args);
if (cmd.check_cmd_line_flag("help")) {
help = true;
return;
}
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("iterations", iterations, 10);
cmd.get_cmd_line_argument("benchmark", benchmark_path);
if (cmd.check_cmd_line_flag("no_verif")) {
verification = false;
}
}
/// Prints the usage statement.
std::ostream & print_usage(std::ostream &out) const {
out << "92_blackwell_moe_gemm_fp4_grouped\n\n"
<< " Blackwell MoE-style grouped NVFP4 GEMM implementation using TMA to load A and CPASYNC to load B\n\n"
<< "Options:\n\n"
<< " --help If specified, displays this usage statement\n\n"
<< " --m=<int> Sets the M extent of the GEMM\n"
<< " --n=<int> Sets the N extent of the GEMM\n"
<< " --k=<int> Sets the K extent of the GEMM\n"
<< " --l=<int> Sets the L extent (batch count) of the GEMM\n"
<< " --iterations=<int> Set the number of profiling iterations to perform\n"
<< " --benchmark=<file> Executes a benchmark problem size\n"
<< " --no_verif Do not run verification kernels\n";
return out;
}
};
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class Element, class Layout>
bool initialize_block(
cutlass::TensorView<Element, Layout> view,
uint64_t seed) {
double scope_max, scope_min;
constexpr int bits_input = cutlass::sizeof_bits<Element>::value;
if constexpr (bits_input == 1) {
scope_max = 2;
scope_min = 0;
}
else if constexpr (bits_input <= 6) {
scope_max = 2;
scope_min = -2;
}
else if constexpr (bits_input <= 8) {
if constexpr (cute::is_same_v<Element, cutlass::float_ue8m0_t> || cute::is_same_v<Element, cutlass::float_ue4m3_t>) {
scope_max = 4;
scope_min = 1;
}
else {
scope_max = 1;
scope_min = -1;
}
}
else{
scope_max = 4;
scope_min = -4;
}
cutlass::reference::host::TensorFillRandomUniform(
view, seed, scope_max, scope_min, 0);
return true;
}
template <class T>
auto make_iterator(T* ptr) {
return cute::recast_ptr<T>(ptr);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
struct ExampleRunner {
// Type of kernel schedule to generate
using MainloopScheduleType = cutlass::gemm::KernelMixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100;
// Type of epilogue schedule to generate
using EpilogueScheduleType = cutlass::epilogue::collective::EpilogueScheduleAuto;
static constexpr bool FuseQuantization = false;
using LayoutATag = cutlass::layout::RowMajor;
using LayoutBTag = cutlass::layout::ColumnMajor;
using LayoutCTag = cutlass::layout::ColumnMajor;
using LayoutDTag = cutlass::layout::ColumnMajor;
using LayoutSFDTag = LayoutDTag; // Layout type for SFD should be same as D matrix operand
using ElementInput = cutlass::float_e2m1_t; // Element type for Input matrix operands
using ElementSF = cutlass::float_ue4m3_t; // Element type for SF matrix operands
using ElementA = cutlass::nv_float4_t<ElementInput>;
using ElementB = cutlass::nv_float4_t<ElementInput>;
using ElementC = cutlass::half_t;
using ElementD = cute::conditional_t<FuseQuantization, ElementInput, ElementC>;
using ElementSFD = ElementSF;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementScalar = float;
using ClusterShapeMNK = Shape<_1,_1,_1>;
using MmaTileMNK = Shape<_128,_64,_256>; // use tile size of N=64 to match real use cases (N is typically very small in decoding stage)
static constexpr int AlignmentA = 32;
static constexpr int AlignmentB = 32;
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
static constexpr int OutputSFVectorSize = 16;
// D = alpha * acc + beta * C
// With BlockScaleFactor generation.
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
OutputSFVectorSize,
ElementD,
ElementCompute,
ElementSFD, LayoutSFDTag,
ElementC>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm100, cutlass::arch::OpClassBlockScaledTensorOp,
MmaTileMNK, ClusterShapeMNK,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutCTag, AlignmentC,
ElementD, LayoutDTag, AlignmentD,
EpilogueScheduleType,
cute::conditional_t<
FuseQuantization,
FusionOperation,
cutlass::epilogue::fusion::LinearCombination<ElementC, ElementAccumulator>>
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm100, cutlass::arch::OpClassBlockScaledTensorOp,
ElementA, LayoutATag, AlignmentA,
ElementB, LayoutBTag, AlignmentB,
ElementAccumulator,
MmaTileMNK, ClusterShapeMNK,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
MainloopScheduleType
>::CollectiveOp;
using ProblemShapeGroup = cutlass::gemm::GroupProblemShape<Shape<int,int,int>>; // <M,N,K> per group
using ProblemShapeMax = Shape<int,int,int,int>; // max <M,N,K,L>
using ProblemShape = cutlass::gemm::MoEProblemShape<ProblemShapeGroup, ProblemShapeMax>;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
ProblemShape,
CollectiveMainloop,
CollectiveEpilogue
>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
using StrideA = typename Gemm::GemmKernel::StrideA;
using LayoutA = decltype(cute::make_layout(make_shape(0,0,0), StrideA{}));
using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride.
using StrideB = typename Gemm::GemmKernel::StrideB;
using LayoutB = decltype(cute::make_layout(make_shape(0,0,0), StrideB{}));
using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride.
using StrideC = typename Gemm::GemmKernel::StrideC;
using LayoutC = decltype(cute::make_layout(make_shape(0,0,0), StrideC{}));
using StrideD = typename Gemm::GemmKernel::StrideD;
using LayoutD = decltype(cute::make_layout(make_shape(0,0,0), StrideD{}));
using FusionOp = typename Gemm::EpilogueOutputOp;
static constexpr bool IsBlockScaleSupported = FusionOp::IsBlockScaleSupported;
using SfdOutputCfg = cutlass::detail::Sm1xxBlockScaledOutputConfig<OutputSFVectorSize>;
using LayoutSFD = typename SfdOutputCfg::LayoutSF;
//
// Data members
//
/// Initialization
StrideA stride_A;
LayoutA layout_A;
LayoutSFA layout_SFA;
StrideB stride_B;
LayoutB layout_B;
LayoutSFB layout_SFB;
StrideC stride_C;
LayoutC layout_C;
StrideD stride_D;
LayoutD layout_D;
LayoutSFD layout_SFD;
uint64_t seed = 0;
cutlass::HostTensor<ElementA::DataType, cutlass::layout::PackedVectorLayout> block_A;
cutlass::HostTensor<ElementA::ScaleFactorType, cutlass::layout::PackedVectorLayout> block_SFA;
cutlass::HostTensor<ElementB::DataType, cutlass::layout::PackedVectorLayout> block_B;
cutlass::HostTensor<ElementB::ScaleFactorType, cutlass::layout::PackedVectorLayout> block_SFB;
cutlass::HostTensor<ElementC, cutlass::layout::PackedVectorLayout> block_C;
cutlass::HostTensor<ElementD, cutlass::layout::PackedVectorLayout> block_D;
cutlass::HostTensor<ElementSFD, cutlass::layout::PackedVectorLayout> block_SFD;
cutlass::HostTensor<ElementD, cutlass::layout::PackedVectorLayout> block_reference_D;
cutlass::HostTensor<ElementSFD, cutlass::layout::PackedVectorLayout> block_reference_SFD;
cutlass::HostTensor<ElementCompute, cutlass::layout::PackedVectorLayout> block_Normconst;
cutlass::DeviceAllocation<typename ProblemShapeGroup::UnderlyingProblemShape> problem_sizes;
//
// Methods
//
bool verify(ProblemShape const& problem_size, float alpha, float beta) {
// Create the arguments for host reference implementation
Tensor tensor_A = make_tensor(make_iterator(block_A.host_data()), layout_A);
Tensor tensor_SFA = make_tensor(block_SFA.host_data(), layout_SFA);
Tensor tensor_B = make_tensor(make_iterator(block_B.host_data()), layout_B);
Tensor tensor_SFB = make_tensor(block_SFB.host_data(), layout_SFB);
// think about how to simplify the gemm3x interface.
cutlass::reference::host::GettBlockScalingMainloopParams<
ElementAccumulator, // ElementAccumulator
decltype(tensor_A), // TensorA
decltype(tensor_SFA), // TensorSfA
decltype(tensor_B), // TensorB
decltype(tensor_SFB) // TensorSfB
> mainloop_params{tensor_A, tensor_SFA, tensor_B, tensor_SFB};
Tensor tensor_C = cute::make_tensor(make_iterator(block_C.host_data()), layout_C);
Tensor tensor_D = cute::make_tensor(make_iterator(block_reference_D.host_data()), layout_D);
Tensor tensor_SFD = make_tensor(block_reference_SFD.host_data(), layout_SFD);
if constexpr (FuseQuantization) {
cutlass::reference::host::GettBlockScalingEpilogueParams<
ElementCompute, // ElementScalar
ElementAccumulator, // ElementAccumulator
ElementCompute, // ElementCompute
decltype(tensor_C), // TensorC
decltype(tensor_D), // TensorD
decltype(tensor_SFD), // TensorSfD
cute::Int<OutputSFVectorSize>,
cutlass::reference::host::SfStrategy::SfDGen
> epilogue_params {alpha, beta, tensor_C, tensor_D, tensor_SFD, block_Normconst.at(cutlass::make_Coord(0))};
cutlass::reference::host::Gemm3x(mainloop_params, epilogue_params);
}
else {
cutlass::reference::host::GettBlockScalingEpilogueParams<
ElementCompute, // ElementScalar
ElementAccumulator, // ElementAccumulator
ElementCompute, // ElementCompute
decltype(tensor_C), // TensorC
decltype(tensor_D) // TensorD
> epilogue_params {alpha, beta, tensor_C, tensor_D };
cutlass::reference::host::Gemm3x(mainloop_params, epilogue_params);
}
bool passed = true;
// 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 [M, N, K] = problem;
// assume all M == maxM
auto refD_view = block_reference_D.host_view().subview(cutlass::make_Coord(M * N), cutlass::make_Coord(i * maxN * maxM));
auto D_view = block_D.host_view().subview(cutlass::make_Coord(M * N), cutlass::make_Coord(i * maxN * maxM));
passed &= cutlass::reference::host::TensorEquals(refD_view, D_view);
}
return passed;
}
/// 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 [M, N, K, L] = problem_shape_MNKL;
// For SFA and SFB tensors layouts
using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
// 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);
layout_D = make_layout(make_shape(M, N, L), stride_D);
layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, L));
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)));
block_D.reset(cutlass::make_Coord(size(layout_D)));
block_reference_D.reset(cutlass::make_Coord(size(layout_D)));
block_reference_SFD.reset(cutlass::make_Coord(size(filter_zeros(layout_SFD))));
block_Normconst.reset(cutlass::make_Coord(1));
block_SFA.reset(cutlass::make_Coord(size(filter_zeros(layout_SFA))));
block_SFB.reset(cutlass::make_Coord(size(filter_zeros(layout_SFB))));
block_SFD.reset(cutlass::make_Coord(size(filter_zeros(layout_SFD))));
initialize_block(block_A.host_view(), seed + 2021);
initialize_block(block_B.host_view(), seed + 2022);
initialize_block(block_C.host_view(), seed + 2023);
initialize_block(block_SFA.host_view(), seed + 2024);
initialize_block(block_SFB.host_view(), seed + 2025);
block_Normconst.at(cutlass::make_Coord(0)) = 2;
block_A.sync_device();
block_B.sync_device();
block_C.sync_device();
block_D.sync_device();
block_SFA.sync_device();
block_SFB.sync_device();
block_SFD.sync_device();
block_Normconst.sync_device();
}
/// Load a benchmark
std::vector<ProblemShapeGroup::UnderlyingProblemShape> benchmark_problems(std::string const& benchmark_path) {
std::vector<ProblemShapeGroup::UnderlyingProblemShape> problem_sizes_host;
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<std::string> 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());
}
problem_sizes_host.push_back({extent.m(), extent.n(), extent.k()});
}
return problem_sizes_host;
}
bool run(Options const& options, cutlass::KernelHardwareInfo const& hw_info) {
auto problem_sizes_host = 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());
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();
initialize(problem_size);
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
},
{ // Epilogue arguments
{},
block_C.device_data(), stride_C,
block_D.device_data(), stride_D
},
hw_info
};
auto f = [&](auto blockscale) {
auto impl = [this](auto& arguments) {
arguments.epilogue.thread.block_scale_factor_ptr = block_SFD.device_data();
arguments.epilogue.thread.norm_constant_ptr = block_Normconst.device_data();
};
if constexpr (decltype(blockscale)::value) {
impl(arguments);
}
};
f(std::bool_constant<IsBlockScaleSupported>());
// arguments.scheduler.max_swizzle_size = options.swizzle;
arguments.epilogue.thread.alpha = 1.0f;
arguments.epilogue.thread.beta = 0.0f;
Gemm gemm_op;
size_t workspace_size = Gemm::get_workspace_size(arguments);
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
cutlass::Status status = gemm_op.can_implement(arguments);
if (status != cutlass::Status::kSuccess) {
std::cerr << "This kernel is not supported. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
status = gemm_op.initialize(arguments, workspace.get());
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to initialize the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
// Run the GEMM
status = gemm_op.run();
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
cudaError_t result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(result) << std::endl;
return false;
}
if (options.verification) {
// Verify that the result is correct
bool passed = verify(problem_size, 1.0f, 0.0f);
std::cout << " Disposition: " << (passed ? "Passed" : "Failed") << std::endl;
if (!passed) {
exit(-1);
return false;
}
}
// Run profiling loop
if (options.iterations > 0)
{
GpuTimer timer;
timer.start();
for (int iter = 0; iter < options.iterations; ++iter) {
CUTLASS_CHECK(gemm_op.initialize(arguments, workspace.get()));
CUTLASS_CHECK(gemm_op.run());
}
timer.stop();
// 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);
std::cout << " Avg runtime : " << avg_runtime_ms << " ms" << std::endl;
std::cout << " TFLOPS : " << flops / 1e12 << std::endl;
}
return true;
}
};
///////////////////////////////////////////////////////////////////////////////////////////////////
int main(int argc, char const **args) {
cudaDeviceProp props;
cudaError_t error = cudaGetDeviceProperties(&props, 0);
if (error != cudaSuccess) {
std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl;
return -1;
}
if (__CUDACC_VER_MAJOR__ < 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 8)) {
std::cerr << "This example requires CUDA 12.8 or newer." << std::endl;
// Returning zero so this test passes on older Toolkits. Its actions are no-op.
return 0;
}
if (!(props.major == 10 && props.minor == 0)) {
std::cerr << "This example requires a GPU of NVIDIA's Blackwell architecture (compute capability 100)." << std::endl;
return 0;
}
//
// Parse options
//
Options options;
options.parse(argc, args);
if (options.help) {
options.print_usage(std::cout) << std::endl;
return 0;
}
if (options.error) {
std::cerr << "Aborting execution." << std::endl;
return -1;
}
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
//
// Run examples
//
// The KernelHardwareInfo struct holds the number of SMs on the GPU with a given device ID. This
// information is used by the underlying kernel.
cutlass::KernelHardwareInfo hw_info;
// Change device_id to another value if you are running on a machine with multiple GPUs and wish
// to use a GPU other than that with device ID 0.
hw_info.device_id = 0;
hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
std::cout << "Running kernel with mixed TMA+CPASYNC load:" << std::endl;
ExampleRunner runner_mixed_tma_cpasync;
runner_mixed_tma_cpasync.run(options, hw_info);
#endif
return 0;
}
examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_fp4_regular.cu
+654
View File
@@ -0,0 +1,654 @@
/***************************************************************************************************
* Copyright (c) 2025 - 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.
*
**************************************************************************************************/
/*! \file
\brief Example of Blackwell MoE-style NVFP4 GEMM implementation using TMA to load A and CPASYNC to load B
This example demonstrates an implementation of GEMM using mixed TMA+CPASYNC to load input matrices.
In the decoding stage of Mixture of Experts (MoE) models, the number of tokens in different experts
can varies a lot, which requires frequently updates of TMA descriptors in TMA-based implementation.
This examples uses CPASYNC to load activation (B) matrix to avoid the overhead of updating TMA descriptors.
This example assumes all experts have the same number of tokens, in which case the GEMM becomes a regular (batched) gemm.
For the realistic use case where each expert may have different number of tokens (grouped GEMM), check 92_blackwell_moe_gemm_fp4_grouped.
Usage:
$ ./examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_fp4_regular
--m=28672 --n=4 --k=4096 --l=8
*/
#include <iostream>
#include "cute/tensor.hpp"
#include "cutlass/cutlass.h"
#include "cutlass/tensor_ref.h"
#include "cutlass/epilogue/collective/default_epilogue.hpp"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/gemm/kernel/tile_scheduler.hpp"
#include "cute/tensor.hpp"
#include "cutlass/tensor_ref.h"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/detail/sm100_blockscaled_layout.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/gemm/kernel/tile_scheduler_params.h"
#include "cutlass/util/command_line.h"
#include "cutlass/util/distribution.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/packed_stride.hpp"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/util/reference/device/gemm.h"
#include "cutlass/util/reference/device/tensor_compare.h"
#include "cutlass/util/reference/host/tensor_fill.h"
#include "cutlass/util/reference/host/gett.hpp"
#include "cutlass/util/reference/host/tensor_norm.h"
#include "cutlass/util/reference/host/tensor_compare.h"
#include "cutlass/util/command_line.h"
#include "cutlass/util/distribution.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/packed_stride.hpp"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/util/reference/device/gemm_complex.h"
#include "cutlass/util/reference/device/tensor_compare.h"
#include "cutlass/util/reference/device/tensor_fill.h"
#include "cutlass/util/reference/device/gemm.h"
#include "cutlass/util/reference/device/tensor_compare.h"
#include "cutlass/util/reference/host/tensor_fill.h"
#include "cutlass/util/reference/host/gett.hpp"
#include "cutlass/util/reference/host/tensor_norm.h"
#include "cutlass/util/reference/host/tensor_compare.h"
#include "helper.h"
using namespace cute;
///////////////////////////////////////////////////////////////////////////////////////////////////
/// Command line options parsing
struct Options {
bool help;
bool error;
bool verification;
int m, n, k, l;
int iterations;
Options():
help(false),
error(false),
verification(true),
m(2048), n(2048), k(2048), l(1),
iterations(10)
{ }
// Parses the command line
void parse(int argc, char const **args) {
cutlass::CommandLine cmd(argc, args);
if (cmd.check_cmd_line_flag("help")) {
help = true;
return;
}
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("iterations", iterations, 10);
if (cmd.check_cmd_line_flag("no_verif")) {
verification = false;
}
}
/// Prints the usage statement.
std::ostream & print_usage(std::ostream &out) const {
out << "92_blackwell_moe_gemm_fp4_regular\n\n"
<< " Blackwell NVFP4 GEMM implementation using TMA to load A and CPASYNC to load B\n\n"
<< "Options:\n\n"
<< " --help If specified, displays this usage statement\n\n"
<< " --m=<int> Sets the M extent of the GEMM\n"
<< " --n=<int> Sets the N extent of the GEMM\n"
<< " --k=<int> Sets the K extent of the GEMM\n"
<< " --l=<int> Sets the L extent (batch count) of the GEMM\n"
<< " --iterations=<int> Set the number of profiling iterations to perform\n"
<< " --no_verif Do not run verification kernels\n";
return out;
}
};
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class Element, class Layout>
bool initialize_block(
cutlass::TensorView<Element, Layout> view,
uint64_t seed) {
double scope_max, scope_min;
constexpr int bits_input = cutlass::sizeof_bits<Element>::value;
if constexpr (bits_input == 1) {
scope_max = 2;
scope_min = 0;
}
else if constexpr (bits_input <= 6) {
scope_max = 2;
scope_min = -2;
}
else if constexpr (bits_input <= 8) {
if constexpr (cute::is_same_v<Element, cutlass::float_ue8m0_t> || cute::is_same_v<Element, cutlass::float_ue4m3_t>) {
scope_max = 4;
scope_min = 1;
}
else {
scope_max = 1;
scope_min = -1;
}
}
else{
scope_max = 4;
scope_min = -4;
}
cutlass::reference::host::TensorFillRandomUniform(
view, seed, scope_max, scope_min, 0);
return true;
}
template <class T>
auto make_iterator(T* ptr) {
return cute::recast_ptr<T>(ptr);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
// MSVC complain about it if moved to ExampleRunner
static constexpr int OutputSFVectorSize = 16;
using SfdOutputCfg = cutlass::detail::Sm1xxBlockScaledOutputConfig<OutputSFVectorSize>;
// Wrapper to construct, run, and verify a GEMM. This example showcases CUTLASS's collective
// operation builders by specializing the GEMM on the kernel+epilogue schedule it will use and the
// number of pipeline stages.
template <
// Type of kernel schedule to generate
class MainloopScheduleType = cutlass::gemm::KernelMixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100,
// Type of epilogue schedule to generate
class EpilogueScheduleType = cutlass::epilogue::collective::EpilogueScheduleAuto,
bool FuseQuantization = false
>
struct ExampleRunner {
using LayoutATag = cutlass::layout::RowMajor;
using LayoutBTag = cutlass::layout::ColumnMajor;
using LayoutCTag = cutlass::layout::ColumnMajor;
using LayoutDTag = cutlass::layout::ColumnMajor;
using LayoutSFDTag = LayoutDTag; // Layout type for SFD should be same as D matrix operand
using ElementInput = cutlass::float_e2m1_t; // Element type for Input matrix operands
using ElementSF = cutlass::float_ue4m3_t; // Element type for SF matrix operands
using ElementA = cutlass::nv_float4_t<ElementInput>;
using ElementB = cutlass::nv_float4_t<ElementInput>;
using ElementC = cutlass::half_t;
using ElementD = cute::conditional_t<FuseQuantization, ElementInput, ElementC>;
using ElementSFD = ElementSF;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementScalar = float;
using ClusterShapeMNK = Shape<_1,_1,_1>;
using MmaTileMNK = Shape<_128,_64,_256>; // use tile size of N=64 to match real use cases (N is typically very small in decoding stage)
static constexpr int AlignmentA = 32;
static constexpr int AlignmentB = 32;
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
// D = alpha * acc + beta * C
// With BlockScaleFactor generation.
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
OutputSFVectorSize,
ElementD,
ElementCompute,
ElementSFD, LayoutSFDTag,
ElementC>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm100, cutlass::arch::OpClassBlockScaledTensorOp,
MmaTileMNK, ClusterShapeMNK,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutCTag, AlignmentC,
ElementD, LayoutDTag, AlignmentD,
EpilogueScheduleType,
cute::conditional_t<
FuseQuantization,
FusionOperation,
cutlass::epilogue::fusion::LinearCombination<ElementC, ElementAccumulator>>
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm100, cutlass::arch::OpClassBlockScaledTensorOp,
ElementA, LayoutATag, AlignmentA,
ElementB, LayoutBTag, AlignmentB,
ElementAccumulator,
MmaTileMNK, ClusterShapeMNK,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
MainloopScheduleType
>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue
>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
using ProblemShapeType = typename Gemm::GemmKernel::ProblemShape;
using StrideA = typename Gemm::GemmKernel::StrideA;
using LayoutA = decltype(cute::make_layout(make_shape(0,0,0), StrideA{}));
using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride.
using StrideB = typename Gemm::GemmKernel::StrideB;
using LayoutB = decltype(cute::make_layout(make_shape(0,0,0), StrideB{}));
using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride.
using StrideC = typename Gemm::GemmKernel::StrideC;
using LayoutC = decltype(cute::make_layout(make_shape(0,0,0), StrideC{}));
using StrideD = typename Gemm::GemmKernel::StrideD;
using LayoutD = decltype(cute::make_layout(make_shape(0,0,0), StrideD{}));
using FusionOp = typename Gemm::EpilogueOutputOp;
static constexpr bool IsBlockScaleSupported = FusionOp::IsBlockScaleSupported;
using LayoutSFD = typename SfdOutputCfg::LayoutSF;
//
// Data members
//
/// Initialization
StrideA stride_A;
LayoutA layout_A;
LayoutSFA layout_SFA;
StrideB stride_B;
LayoutB layout_B;
LayoutSFB layout_SFB;
StrideC stride_C;
LayoutC layout_C;
StrideD stride_D;
LayoutD layout_D;
LayoutSFD layout_SFD;
uint64_t seed = 0;
cutlass::HostTensor<ElementA::DataType, cutlass::layout::PackedVectorLayout> block_A;
cutlass::HostTensor<ElementA::ScaleFactorType, cutlass::layout::PackedVectorLayout> block_SFA;
cutlass::HostTensor<ElementB::DataType, cutlass::layout::PackedVectorLayout> block_B;
cutlass::HostTensor<ElementB::ScaleFactorType, cutlass::layout::PackedVectorLayout> block_SFB;
cutlass::HostTensor<ElementC, cutlass::layout::PackedVectorLayout> block_C;
cutlass::HostTensor<ElementD, cutlass::layout::PackedVectorLayout> block_D;
cutlass::HostTensor<ElementSFD, cutlass::layout::PackedVectorLayout> block_SFD;
cutlass::HostTensor<ElementD, cutlass::layout::PackedVectorLayout> block_reference_D;
cutlass::HostTensor<ElementSFD, cutlass::layout::PackedVectorLayout> block_reference_SFD;
cutlass::HostTensor<ElementCompute, cutlass::layout::PackedVectorLayout> block_Normconst;
//
// Methods
//
bool verify(ProblemShapeType const& problem_size, float alpha, float beta) {
// Create the arguments for host reference implementation
Tensor tensor_A = make_tensor(make_iterator(block_A.host_data()), layout_A);
Tensor tensor_SFA = make_tensor(block_SFA.host_data(), layout_SFA);
Tensor tensor_B = make_tensor(make_iterator(block_B.host_data()), layout_B);
Tensor tensor_SFB = make_tensor(block_SFB.host_data(), layout_SFB);
// think about how to simplify the gemm3x interface.
cutlass::reference::host::GettBlockScalingMainloopParams<
ElementAccumulator, // ElementAccumulator
decltype(tensor_A), // TensorA
decltype(tensor_SFA), // TensorSfA
decltype(tensor_B), // TensorB
decltype(tensor_SFB) // TensorSfB
> mainloop_params{tensor_A, tensor_SFA, tensor_B, tensor_SFB};
Tensor tensor_C = cute::make_tensor(make_iterator(block_C.host_data()), layout_C);
Tensor tensor_D = cute::make_tensor(make_iterator(block_reference_D.host_data()), layout_D);
Tensor tensor_SFD = make_tensor(block_reference_SFD.host_data(), layout_SFD);
if constexpr (FuseQuantization) {
cutlass::reference::host::GettBlockScalingEpilogueParams<
ElementCompute, // ElementScalar
ElementAccumulator, // ElementAccumulator
ElementCompute, // ElementCompute
decltype(tensor_C), // TensorC
decltype(tensor_D), // TensorD
decltype(tensor_SFD), // TensorSfD
cute::Int<OutputSFVectorSize>,
cutlass::reference::host::SfStrategy::SfDGen
> epilogue_params {alpha, beta, tensor_C, tensor_D, tensor_SFD, block_Normconst.at(cutlass::make_Coord(0))};
cutlass::reference::host::Gemm3x(mainloop_params, epilogue_params);
}
else {
cutlass::reference::host::GettBlockScalingEpilogueParams<
ElementCompute, // ElementScalar
ElementAccumulator, // ElementAccumulator
ElementCompute, // ElementCompute
decltype(tensor_C), // TensorC
decltype(tensor_D) // TensorD
> epilogue_params {alpha, beta, tensor_C, tensor_D };
cutlass::reference::host::Gemm3x(mainloop_params, epilogue_params);
}
bool passed = true, passed_sfd = true;
// Comparison
block_D.sync_host();
passed &= cutlass::reference::host::TensorEquals(block_reference_D.host_view(), block_D.host_view());
if constexpr (FuseQuantization) {
passed &= (cutlass::reference::host::TensorNorm(block_reference_D.host_view()) > 0);
passed &= (cutlass::reference::host::TensorNorm(block_D.host_view()) > 0);
block_SFD.sync_host();
passed_sfd &= cutlass::reference::host::TensorEquals(block_reference_SFD.host_view(), block_SFD.host_view());
passed_sfd &= (cutlass::reference::host::TensorNorm(block_reference_SFD.host_view()) > 0);
passed_sfd &= (cutlass::reference::host::TensorNorm(block_SFD.host_view()) > 0);
}
// printf("passed=%d\n", (int)passed);
// printf("passed_sfd=%d\n", (int)passed_sfd);
return passed && passed_sfd;
}
/// Initialize operands to be used in the GEMM and reference GEMM
void initialize(ProblemShapeType const& problem_size) {
auto problem_shape_MNKL = cute::append<4>(problem_size, 1);
auto [M, N, K, L] = problem_shape_MNKL;
// For SFA and SFB tensors layouts
using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
// 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);
layout_D = make_layout(make_shape(M, N, L), stride_D);
layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, L));
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)));
block_D.reset(cutlass::make_Coord(size(layout_D)));
block_reference_D.reset(cutlass::make_Coord(size(layout_D)));
block_reference_SFD.reset(cutlass::make_Coord(size(filter_zeros(layout_SFD))));
block_Normconst.reset(cutlass::make_Coord(1));
block_SFA.reset(cutlass::make_Coord(size(filter_zeros(layout_SFA))));
block_SFB.reset(cutlass::make_Coord(size(filter_zeros(layout_SFB))));
block_SFD.reset(cutlass::make_Coord(size(filter_zeros(layout_SFD))));
initialize_block(block_A.host_view(), seed + 2021);
initialize_block(block_B.host_view(), seed + 2022);
initialize_block(block_C.host_view(), seed + 2023);
initialize_block(block_SFA.host_view(), seed + 2024);
initialize_block(block_SFB.host_view(), seed + 2025);
block_Normconst.at(cutlass::make_Coord(0)) = 2;
block_A.sync_device();
block_B.sync_device();
block_C.sync_device();
block_D.sync_device();
block_SFA.sync_device();
block_SFB.sync_device();
block_SFD.sync_device();
block_Normconst.sync_device();
}
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
};
if constexpr (IsBlockScaleSupported) {
arguments.epilogue.thread.block_scale_factor_ptr = block_SFD.device_data();
arguments.epilogue.thread.norm_constant_ptr = block_Normconst.device_data();
}
// arguments.scheduler.max_swizzle_size = options.swizzle;
arguments.epilogue.thread.alpha = 1.0f;
arguments.epilogue.thread.beta = 0.0f;
Gemm gemm_op;
size_t workspace_size = Gemm::get_workspace_size(arguments);
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
cutlass::Status status = gemm_op.can_implement(arguments);
if (status != cutlass::Status::kSuccess) {
std::cerr << "This kernel is not supported. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
status = gemm_op.initialize(arguments, workspace.get());
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to initialize the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
// Run the GEMM
status = gemm_op.run();
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
cudaError_t result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(result) << std::endl;
return false;
}
if (options.verification) {
// Verify that the result is correct
bool passed = verify(problem_size, 1.0f, 0.0f);
std::cout << " Disposition: " << (passed ? "Passed" : "Failed") << std::endl;
if (!passed) {
exit(-1);
return false;
}
}
// Run profiling loop
if (options.iterations > 0)
{
GpuTimer timer;
timer.start();
for (int iter = 0; iter < options.iterations; ++iter) {
CUTLASS_CHECK(gemm_op.initialize(arguments, workspace.get()));
CUTLASS_CHECK(gemm_op.run());
}
timer.stop();
// 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);
std::cout << " Avg runtime : " << avg_runtime_ms << " ms" << std::endl;
std::cout << " TFLOPS : " << flops / 1e12 << std::endl;
}
return true;
}
};
///////////////////////////////////////////////////////////////////////////////////////////////////
int main(int argc, char const **args) {
cudaDeviceProp props;
cudaError_t error = cudaGetDeviceProperties(&props, 0);
if (error != cudaSuccess) {
std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl;
return -1;
}
if (__CUDACC_VER_MAJOR__ < 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 8)) {
std::cerr << "This example requires CUDA 12.8 or newer." << std::endl;
// Returning zero so this test passes on older Toolkits. Its actions are no-op.
return 0;
}
if (!(props.major == 10 && props.minor == 0)) {
std::cerr << "This example requires a GPU of NVIDIA's Blackwell architecture (compute capability 100)." << std::endl;
return 0;
}
//
// Parse options
//
Options options;
options.parse(argc, args);
if (options.help) {
options.print_usage(std::cout) << std::endl;
return 0;
}
if (options.error) {
std::cerr << "Aborting execution." << std::endl;
return -1;
}
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
//
// Run examples
//
// The KernelHardwareInfo struct holds the number of SMs on the GPU with a given device ID. This
// information is used by the underlying kernel.
cutlass::KernelHardwareInfo hw_info;
// Change device_id to another value if you are running on a machine with multiple GPUs and wish
// to use a GPU other than that with device ID 0.
hw_info.device_id = 0;
hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
std::cout << "Running kernel with TMA load:" << std::endl;
ExampleRunner<cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledSm100> runner_tma;
runner_tma.run(options, hw_info);
std::cout << "Running kernel with mixed TMA+CPASYNC load:" << std::endl;
ExampleRunner<cutlass::gemm::KernelMixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100> runner_mixed_tma_cpasync;
runner_mixed_tma_cpasync.run(options, hw_info);
#endif
return 0;
}
examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_grouped.cu
+541
View File
@@ -0,0 +1,541 @@
/***************************************************************************************************
* Copyright (c) 2025 - 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.
*
**************************************************************************************************/
/*! \file
\brief Example of Blackwell MoE-style grouped GEMM implementation using TMA to load A and CPASYNC to load B.
This example demonstrates an implementation of GEMM using mixed TMA+CPASYNC to load input matrices.
In the decoding stage of Mixture of Experts (MoE) models, the number of tokens in different experts
can varies a lot, which requires frequently updates of TMA descriptors in TMA-based implementation.
This examples uses CPASYNC to load activation (B) matrix to avoid the overhead of updating TMA descriptors.
Usage:
$ ./examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_grouped
--m=28672 --n=4 --k=4096 --l=8 --benchmark=benchmark.txt
*/
#include <iostream>
#include <fstream>
#include "cute/tensor.hpp"
#include "cutlass/cutlass.h"
#include "cutlass/tensor_ref.h"
#include "cutlass/epilogue/collective/default_epilogue.hpp"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/gemm/kernel/tile_scheduler.hpp"
#include "cutlass/util/command_line.h"
#include "cutlass/util/distribution.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/packed_stride.hpp"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/util/reference/device/gemm.h"
#include "cutlass/util/reference/device/tensor_compare.h"
#include "cutlass/util/reference/device/tensor_fill.h"
#include "helper.h"
using namespace cute;
///////////////////////////////////////////////////////////////////////////////////////////////////
/// Command line options parsing
struct Options {
bool help;
bool error;
bool verification;
int m, n, k, l;
int iterations;
std::string benchmark_path;
Options():
help(false),
error(false),
verification(true),
m(2048), n(2048), k(2048), l(1),
iterations(10)
{ }
// Parses the command line
void parse(int argc, char const **args) {
cutlass::CommandLine cmd(argc, args);
if (cmd.check_cmd_line_flag("help")) {
help = true;
return;
}
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("iterations", iterations, 10);
cmd.get_cmd_line_argument("benchmark", benchmark_path);
if (cmd.check_cmd_line_flag("no_verif")) {
verification = false;
}
}
/// Prints the usage statement.
std::ostream & print_usage(std::ostream &out) const {
out << "92_blackwell_moe_gemm_grouped\n\n"
<< " Blackwell MoE-style grouped GEMM implementation using TMA to load A and CPASYNC to load B\n\n"
<< "Options:\n\n"
<< " --help If specified, displays this usage statement\n\n"
<< " --m=<int> Sets the M extent of the GEMM\n"
<< " --n=<int> Sets the N extent of the GEMM\n"
<< " --k=<int> Sets the K extent of the GEMM\n"
<< " --l=<int> Sets the L extent (batch count) of the GEMM\n"
<< " --iterations=<int> Set the number of profiling iterations to perform\n"
<< " --benchmark=<file> Executes a benchmark problem size\n"
<< " --no_verif Do not run verification kernels\n";
return out;
}
};
///////////////////////////////////////////////////////////////////////////////////////////////////
/// Helper to initialize a block of device data
template <class Element>
bool initialize_block(
cutlass::DeviceAllocation<Element>& block,
uint64_t seed=2023) {
Element scope_max, scope_min;
int bits_input = cutlass::sizeof_bits<Element>::value;
if (bits_input == 1) {
scope_max = static_cast<Element>(2);
scope_min = static_cast<Element>(0);
}
else if (bits_input <= 8) {
scope_max = static_cast<Element>(2);
scope_min = static_cast<Element>(-2);
}
else {
scope_max = static_cast<Element>(8);
scope_min = static_cast<Element>(-8);
}
cutlass::reference::device::BlockFillRandomUniform(
block.get(), block.size(), seed, scope_max, scope_min, 0);
return true;
}
///////////////////////////////////////////////////////////////////////////////////////////////////
struct ExampleRunner {
// Type of kernel schedule to generate
using MainloopScheduleType = cutlass::gemm::KernelMixedTmaCpAsyncWarpSpecialized1SmSm100;
// Type of epilogue schedule to generate
using EpilogueScheduleType = cutlass::epilogue::collective::EpilogueScheduleAuto;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::ColumnMajor;
using ElementA = cutlass::float_e4m3_t;
using ElementB = cutlass::float_e4m3_t;
using ElementC = cutlass::half_t;
using ElementD = cutlass::half_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementScalar = float;
using ClusterShapeMNK = Shape<_1,_1,_1>;
using MmaTileMNK = Shape<_128,_16,Int<128 / sizeof(ElementA)>>; // use tile size of N=16 to match real use cases (N is typically very small in decoding stage)
// 16B alignment lets us use TMA
static constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementA>::value;
static constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value;
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
MmaTileMNK, ClusterShapeMNK,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
EpilogueScheduleType,
cutlass::epilogue::fusion::LinearCombination<ElementC, ElementAccumulator>
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
ElementA, LayoutA, AlignmentA,
ElementB, LayoutB, AlignmentB,
ElementAccumulator,
MmaTileMNK, ClusterShapeMNK,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
MainloopScheduleType
>::CollectiveOp;
using ProblemShapeGroup = cutlass::gemm::GroupProblemShape<Shape<int,int,int>>; // <M,N,K> per group
using ProblemShapeMax = Shape<int,int,int,int>; // max <M,N,K,L>
using ProblemShape = cutlass::gemm::MoEProblemShape<ProblemShapeGroup, ProblemShapeMax>;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
ProblemShape,
CollectiveMainloop,
CollectiveEpilogue
//, cutlass::gemm::MoEScheduler
>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
// using ProblemShapeType = typename Gemm::GemmKernel::ProblemShape;
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 LayoutTagA = cutlass::gemm::detail::StrideToLayoutTagA_t<StrideA>;
using LayoutTagB = cutlass::gemm::detail::StrideToLayoutTagB_t<StrideB>;
using LayoutTagC = cutlass::gemm::detail::StrideToLayoutTagC_t<StrideC>;
using LayoutTagD = cutlass::gemm::detail::StrideToLayoutTagC_t<StrideD>;
//
// Data members
//
/// Initialization
StrideA stride_A;
StrideB stride_B;
StrideC stride_C;
StrideD stride_D;
uint64_t seed = 0;
cutlass::DeviceAllocation<typename Gemm::ElementA> block_A;
cutlass::DeviceAllocation<typename Gemm::ElementB> block_B;
cutlass::DeviceAllocation<typename Gemm::ElementC> block_C;
cutlass::DeviceAllocation<typename Gemm::ElementD> block_D;
cutlass::DeviceAllocation<typename Gemm::ElementD> block_ref_D;
cutlass::DeviceAllocation<typename ProblemShapeGroup::UnderlyingProblemShape> problem_sizes;
//
// Methods
//
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 [M, N, K] = problem;
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));
cutlass::TensorRef ref_C(block_C.get() + size_t(1) * i * maxN * maxM, Gemm::LayoutC(maxM));
cutlass::TensorRef ref_D(block_ref_D.get() + size_t(1) * i * maxN * maxM, Gemm::LayoutD(maxM));
using DeviceGemmReference = cutlass::reference::device::Gemm<
ElementA,
LayoutA,
ElementB,
LayoutB,
ElementC,
LayoutC,
ElementScalar,
ElementAccumulator>;
DeviceGemmReference gemm_reference;
gemm_reference(
{M, N, K},
ElementScalar(alpha),
ref_A,
ref_B,
ElementScalar(beta),
ref_C,
ref_D);
cudaError_t result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
std::cerr << "Reference kernel failed. Last CUDA error: "
<< cudaGetErrorString(result) << std::endl;
return false;
}
// Check if output from CUTLASS kernel and reference kernel are equal or not
// assume all M == maxM
bool passed = cutlass::reference::device::BlockCompareEqual(block_ref_D.get() + size_t(1) * i * maxN * maxM, block_D.get() + size_t(1) * i * maxN * maxM, M * N);
if (!passed) {
return false;
}
}
return true;
}
/// 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 [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));
block_A.reset(size_t(1) * M * K * L);
block_B.reset(size_t(1) * K * N * L);
block_C.reset(size_t(1) * M * N * L);
block_D.reset(size_t(1) * M * N * L);
block_ref_D.reset(size_t(1) * M * N * L);
initialize_block(block_A, seed + 2023);
initialize_block(block_B, seed + 2022);
initialize_block(block_C, seed + 2021);
}
/// Load a benchmark
std::vector<ProblemShapeGroup::UnderlyingProblemShape> benchmark_problems(std::string const& benchmark_path) {
std::vector<ProblemShapeGroup::UnderlyingProblemShape> problem_sizes_host;
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<std::string> 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());
}
problem_sizes_host.push_back({extent.m(), extent.n(), extent.k()});
}
return problem_sizes_host;
}
bool run(Options const& options, cutlass::KernelHardwareInfo const& hw_info) {
auto problem_sizes_host = 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());
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();
initialize(problem_size);
typename Gemm::Arguments arguments{
cutlass::gemm::GemmUniversalMode::kGrouped,
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
};
// arguments.scheduler.max_swizzle_size = options.swizzle;
arguments.epilogue.thread.alpha = 1.0f;
arguments.epilogue.thread.beta = 0.0f;
Gemm gemm_op;
size_t workspace_size = Gemm::get_workspace_size(arguments);
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
cutlass::Status status = gemm_op.can_implement(arguments);
if (status != cutlass::Status::kSuccess) {
std::cerr << "This kernel is not supported. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
status = gemm_op.initialize(arguments, workspace.get());
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to initialize the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
// Run the GEMM
status = gemm_op.run();
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
cudaError_t result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(result) << std::endl;
return false;
}
if (options.verification) {
// Verify that the result is correct
bool passed = verify(problem_size, 1.0f, 0.0f);
std::cout << " Disposition: " << (passed ? "Passed" : "Failed") << std::endl;
if (!passed) {
exit(-1);
return false;
}
}
// Run profiling loop
if (options.iterations > 0)
{
GpuTimer timer;
timer.start();
for (int iter = 0; iter < options.iterations; ++iter) {
CUTLASS_CHECK(gemm_op.initialize(arguments, workspace.get()));
CUTLASS_CHECK(gemm_op.run());
}
timer.stop();
// 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);
std::cout << " Avg runtime : " << avg_runtime_ms << " ms" << std::endl;
std::cout << " TFLOPS : " << flops / 1e12 << std::endl;
}
return true;
}
};
///////////////////////////////////////////////////////////////////////////////////////////////////
int main(int argc, char const **args) {
cudaDeviceProp props;
cudaError_t error = cudaGetDeviceProperties(&props, 0);
if (error != cudaSuccess) {
std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl;
return -1;
}
if (__CUDACC_VER_MAJOR__ < 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 8)) {
std::cerr << "This example requires CUDA 12.8 or newer." << std::endl;
// Returning zero so this test passes on older Toolkits. Its actions are no-op.
return 0;
}
if (!(props.major == 10 && props.minor == 0)) {
std::cerr << "This example requires a GPU of NVIDIA's Blackwell architecture (compute capability 100)." << std::endl;
return 0;
}
//
// Parse options
//
Options options;
options.parse(argc, args);
if (options.help) {
options.print_usage(std::cout) << std::endl;
return 0;
}
if (options.error) {
std::cerr << "Aborting execution." << std::endl;
return -1;
}
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
//
// Run examples
//
// The KernelHardwareInfo struct holds the number of SMs on the GPU with a given device ID. This
// information is used by the underlying kernel.
cutlass::KernelHardwareInfo hw_info;
// Change device_id to another value if you are running on a machine with multiple GPUs and wish
// to use a GPU other than that with device ID 0.
hw_info.device_id = 0;
hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
std::cout << "Running kernel with mixed TMA+CPASYNC load:" << std::endl;
ExampleRunner runner_mixed_tma_cpasync;
runner_mixed_tma_cpasync.run(options, hw_info);
#endif
return 0;
}
examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_regular.cu
+484
View File
@@ -0,0 +1,484 @@
/***************************************************************************************************
* Copyright (c) 2025 - 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.
*
**************************************************************************************************/
/*! \file
\brief Example of Blackwell MoE-style GEMM implementation using TMA to load A and CPASYNC to load B.
This example demonstrates an implementation of GEMM using mixed TMA+CPASYNC to load input matrices.
In the decoding stage of Mixture of Experts (MoE) models, the number of tokens in different experts
can varies a lot, which requires frequently updates of TMA descriptors in TMA-based implementation.
This examples uses CPASYNC to load activation (B) matrix to avoid the overhead of updating TMA descriptors.
Usage:
$ ./examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_regular
--m=28672 --n=4 --k=4096 --l=8
*/
#include <iostream>
#include "cute/tensor.hpp"
#include "cutlass/cutlass.h"
#include "cutlass/tensor_ref.h"
#include "cutlass/epilogue/collective/default_epilogue.hpp"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/gemm/kernel/tile_scheduler.hpp"
#include "cutlass/util/command_line.h"
#include "cutlass/util/distribution.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/packed_stride.hpp"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/util/reference/device/gemm_complex.h"
#include "cutlass/util/reference/device/tensor_compare.h"
#include "cutlass/util/reference/device/tensor_fill.h"
#include "helper.h"
using namespace cute;
///////////////////////////////////////////////////////////////////////////////////////////////////
/// Command line options parsing
struct Options {
bool help;
bool error;
bool verification;
int m, n, k, l;
int iterations;
Options():
help(false),
error(false),
verification(true),
m(2048), n(2048), k(2048), l(1),
iterations(10)
{ }
// Parses the command line
void parse(int argc, char const **args) {
cutlass::CommandLine cmd(argc, args);
if (cmd.check_cmd_line_flag("help")) {
help = true;
return;
}
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("iterations", iterations, 10);
if (cmd.check_cmd_line_flag("no_verif")) {
verification = false;
}
}
/// Prints the usage statement.
std::ostream & print_usage(std::ostream &out) const {
out << "92_blackwell_moe_gemm_regular\n\n"
<< " Blackwell GEMM implementation using TMA to load A and CPASYNC to load B\n\n"
<< "Options:\n\n"
<< " --help If specified, displays this usage statement\n\n"
<< " --m=<int> Sets the M extent of the GEMM\n"
<< " --n=<int> Sets the N extent of the GEMM\n"
<< " --k=<int> Sets the K extent of the GEMM\n"
<< " --l=<int> Sets the L extent (batch count) of the GEMM\n"
<< " --iterations=<int> Set the number of profiling iterations to perform\n"
<< " --no_verif Do not run verification kernels\n";
return out;
}
};
///////////////////////////////////////////////////////////////////////////////////////////////////
/// Helper to initialize a block of device data
template <class Element>
bool initialize_block(
cutlass::DeviceAllocation<Element>& block,
uint64_t seed=2023) {
Element scope_max, scope_min;
int bits_input = cutlass::sizeof_bits<Element>::value;
if (bits_input == 1) {
scope_max = static_cast<Element>(2);
scope_min = static_cast<Element>(0);
}
else if (bits_input <= 8) {
scope_max = static_cast<Element>(2);
scope_min = static_cast<Element>(-2);
}
else {
scope_max = static_cast<Element>(8);
scope_min = static_cast<Element>(-8);
}
cutlass::reference::device::BlockFillRandomUniform(
block.get(), block.size(), seed, scope_max, scope_min, 0);
return true;
}
///////////////////////////////////////////////////////////////////////////////////////////////////
// Wrapper to construct, run, and verify a GEMM. This example showcases CUTLASS's collective
// operation builders by specializing the GEMM on the kernel+epilogue schedule it will use and the
// number of pipeline stages.
template <
// Type of kernel schedule to generate
class MainloopScheduleType = cutlass::gemm::KernelMixedTmaCpAsyncWarpSpecialized1SmSm100,
// Type of epilogue schedule to generate
class EpilogueScheduleType = cutlass::epilogue::collective::EpilogueScheduleAuto
>
struct ExampleRunner {
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::ColumnMajor;
using ElementA = cutlass::float_e4m3_t;
using ElementB = cutlass::float_e4m3_t;
using ElementC = cutlass::half_t;
using ElementD = cutlass::half_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementScalar = float;
using ClusterShapeMNK = Shape<_1,_1,_1>;
using MmaTileMNK = Shape<_128,_16,Int<128 / sizeof(ElementA)>>; // use tile size of N=16 to match real use cases (N is typically very small in decoding stage)
// 16B alignment lets us use TMA
static constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementA>::value;
static constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value;
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
MmaTileMNK, ClusterShapeMNK,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
EpilogueScheduleType,
cutlass::epilogue::fusion::LinearCombination<ElementC, ElementAccumulator>
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
ElementA, LayoutA, AlignmentA,
ElementB, LayoutB, AlignmentB,
ElementAccumulator,
MmaTileMNK, ClusterShapeMNK,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
MainloopScheduleType
>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue
>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
using ProblemShapeType = typename Gemm::GemmKernel::ProblemShape;
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 LayoutTagA = cutlass::gemm::detail::StrideToLayoutTagA_t<StrideA>;
using LayoutTagB = cutlass::gemm::detail::StrideToLayoutTagB_t<StrideB>;
using LayoutTagC = cutlass::gemm::detail::StrideToLayoutTagC_t<StrideC>;
using LayoutTagD = cutlass::gemm::detail::StrideToLayoutTagC_t<StrideD>;
//
// Data members
//
/// Initialization
StrideA stride_A;
StrideB stride_B;
StrideC stride_C;
StrideD stride_D;
uint64_t seed = 0;
cutlass::DeviceAllocation<typename Gemm::ElementA> block_A;
cutlass::DeviceAllocation<typename Gemm::ElementB> block_B;
cutlass::DeviceAllocation<typename Gemm::ElementC> block_C;
cutlass::DeviceAllocation<typename Gemm::ElementD> block_D;
cutlass::DeviceAllocation<typename Gemm::ElementD> block_ref_D;
//
// Methods
//
bool verify(ProblemShapeType const& problem_size, float alpha, float beta) {
auto [M, N, K, L] = problem_size;
cutlass::TensorRef ref_A(block_A.get(), Gemm::LayoutA::packed({M, K}));
cutlass::TensorRef ref_B(block_B.get(), Gemm::LayoutB::packed({K, N}));
cutlass::TensorRef ref_C(block_C.get(), Gemm::LayoutC::packed({M, N}));
cutlass::TensorRef ref_D(block_ref_D.get(), Gemm::LayoutD::packed({M, N}));
cutlass::reference::device::GemmComplex(
{M, N, K},
ElementScalar(alpha),
ref_A,
cutlass::ComplexTransform::kNone,
ref_B,
cutlass::ComplexTransform::kNone,
ElementScalar(beta),
ref_C,
ref_D,
ElementAccumulator(0),
L, // batch_count
M * K, // batch_stride_A
K * N, // batch_stride_B
M * N, // batch_stride_C
M * N // batch_stride_D
);
cudaError_t result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
std::cerr << "Reference kernel failed. Last CUDA error: "
<< cudaGetErrorString(result) << std::endl;
return false;
}
// Check if output from CUTLASS kernel and reference kernel are equal or not
bool passed = cutlass::reference::device::BlockCompareEqual(block_ref_D.get(), block_D.get(), block_D.size());
return passed;
}
/// Initialize operands to be used in the GEMM and reference GEMM
void initialize(ProblemShapeType const& problem_size) {
auto problem_shape_MNKL = cute::append<4>(problem_size, 1);
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));
block_A.reset(size_t(1) * M * K * L);
block_B.reset(size_t(1) * K * N * L);
block_C.reset(size_t(1) * M * N * L);
block_D.reset(size_t(1) * M * N * L);
block_ref_D.reset(size_t(1) * M * N * L);
initialize_block(block_A, seed + 2023);
initialize_block(block_B, seed + 2022);
initialize_block(block_C, seed + 2021);
}
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
};
// arguments.scheduler.max_swizzle_size = options.swizzle;
arguments.epilogue.thread.alpha = 1.0f;
arguments.epilogue.thread.beta = 0.0f;
Gemm gemm_op;
size_t workspace_size = Gemm::get_workspace_size(arguments);
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
cutlass::Status status = gemm_op.can_implement(arguments);
if (status != cutlass::Status::kSuccess) {
std::cerr << "This kernel is not supported. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
status = gemm_op.initialize(arguments, workspace.get());
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to initialize the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
// Run the GEMM
status = gemm_op.run();
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
cudaError_t result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(result) << std::endl;
return false;
}
if (options.verification) {
// Verify that the result is correct
bool passed = verify(problem_size, 1.0f, 0.0f);
std::cout << " Disposition: " << (passed ? "Passed" : "Failed") << std::endl;
if (!passed) {
exit(-1);
return false;
}
}
// Run profiling loop
if (options.iterations > 0)
{
GpuTimer timer;
timer.start();
for (int iter = 0; iter < options.iterations; ++iter) {
CUTLASS_CHECK(gemm_op.initialize(arguments, workspace.get()));
CUTLASS_CHECK(gemm_op.run());
}
timer.stop();
// 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);
std::cout << " Avg runtime : " << avg_runtime_ms << " ms" << std::endl;
std::cout << " TFLOPS : " << flops / 1e12 << std::endl;
}
return true;
}
};
///////////////////////////////////////////////////////////////////////////////////////////////////
int main(int argc, char const **args) {
cudaDeviceProp props;
cudaError_t error = cudaGetDeviceProperties(&props, 0);
if (error != cudaSuccess) {
std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl;
return -1;
}
if (__CUDACC_VER_MAJOR__ < 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 8)) {
std::cerr << "This example requires CUDA 12.8 or newer." << std::endl;
// Returning zero so this test passes on older Toolkits. Its actions are no-op.
return 0;
}
if (!(props.major == 10 && props.minor == 0)) {
std::cerr << "This example requires a GPU of NVIDIA's Blackwell architecture (compute capability 100)." << std::endl;
return 0;
}
//
// Parse options
//
Options options;
options.parse(argc, args);
if (options.help) {
options.print_usage(std::cout) << std::endl;
return 0;
}
if (options.error) {
std::cerr << "Aborting execution." << std::endl;
return -1;
}
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
//
// Run examples
//
// The KernelHardwareInfo struct holds the number of SMs on the GPU with a given device ID. This
// information is used by the underlying kernel.
cutlass::KernelHardwareInfo hw_info;
// Change device_id to another value if you are running on a machine with multiple GPUs and wish
// to use a GPU other than that with device ID 0.
hw_info.device_id = 0;
hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
std::cout << "Running kernel with TMA load:" << std::endl;
ExampleRunner<cutlass::gemm::KernelTmaWarpSpecialized1SmSm100> runner_tma;
runner_tma.run(options, hw_info);
std::cout << "Running kernel with CPASYNC load:" << std::endl;
ExampleRunner<cutlass::gemm::KernelWarpSpecialized1SmSm100> runner_cpasync;
runner_cpasync.run(options, hw_info);
std::cout << "Running kernel with mixed TMA+CPASYNC load:" << std::endl;
ExampleRunner<cutlass::gemm::KernelMixedTmaCpAsyncWarpSpecialized1SmSm100> runner_mixed_tma_cpasync;
runner_mixed_tma_cpasync.run(options, hw_info);
#endif
return 0;
}
examples/92_blackwell_moe_gemm/CMakeLists.txt
+70
View File
@@ -0,0 +1,70 @@
# Copyright (c) 2025 - 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.
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_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)
if (CUTLASS_NVCC_ARCHS MATCHES 100a)
cutlass_example_add_executable(
92_blackwell_moe_gemm_regular
92_blackwell_moe_gemm_regular.cu
TEST_COMMAND_OPTIONS
TEST_MIXTRAL_A
TEST_MIXTRAL_B
TEST_DEEPSEEK_A
TEST_DEEPSEEK_B
TEST_IRREGULAR_MNK
)
cutlass_example_add_executable(
92_blackwell_moe_gemm_grouped
92_blackwell_moe_gemm_grouped.cu
)
cutlass_example_add_executable(
92_blackwell_moe_gemm_fp4_regular
92_blackwell_moe_gemm_fp4_regular.cu
TEST_COMMAND_OPTIONS
TEST_MIXTRAL_A
TEST_MIXTRAL_B
TEST_DEEPSEEK_A_FP4
TEST_DEEPSEEK_B_FP4
TEST_IRREGULAR_MNK_FP4
)
cutlass_example_add_executable(
92_blackwell_moe_gemm_fp4_grouped
92_blackwell_moe_gemm_fp4_grouped.cu
)
endif()
examples/CMakeLists.txt
+1
View File
@@ -169,6 +169,7 @@ foreach(EXAMPLE
89_sm103_fp4_ultra_gemm
90_sm103_fp4_ultra_grouped_gemm
91_fp4_gemv
92_blackwell_moe_gemm
)
add_subdirectory(${EXAMPLE})
examples/python/CuTeDSL/ampere/all_reduce.py
+314
View File
@@ -0,0 +1,314 @@
import os
import torch
import argparse
from cuda import cuda
from cuda.bindings import driver
from typing import Type
import torch.distributed as dist
import torch.distributed._symmetric_memory as symm_mem
import torch.multiprocessing as mp
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
from cutlass._mlir.dialects import llvm, builtin, vector, arith
WORLD_SIZE = 8
PING_PONG_SIZE = 3
def setup(rank, world_size):
# set environment variables for torch.distributed environment
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "12959"
dist.init_process_group("nccl", rank=rank, world_size=world_size)
torch.cuda.set_device(rank)
def cleanup():
dist.destroy_process_group()
class AllReduceKernel:
@cute.jit
def __call__(
self,
rank,
signal,
local_input: cute.Tensor,
local_output: cute.Tensor,
buffer0: cute.Tensor,
buffer1: cute.Tensor,
buffer2: cute.Tensor,
buffer3: cute.Tensor,
buffer4: cute.Tensor,
buffer5: cute.Tensor,
buffer6: cute.Tensor,
buffer7: cute.Tensor,
stream: cuda.CUstream,
):
# define constants for future use
num_of_elements = cute.size(local_input.layout)
# 128 threads per block and 4 elements per thread
tv_layout = cute.make_layout(((128), (4)), stride=((1), (1)))
tile = cute.size(tv_layout.shape)
buffers = [
buffer0,
buffer1,
buffer2,
buffer3,
buffer4,
buffer5,
buffer6,
buffer7,
]
tiled_buffers = [
cute.logical_divide(buffer, (tile, None, None)) for buffer in buffers
]
tiled_input = cute.zipped_divide(local_input, cute.make_layout(tile))
tiled_output = cute.zipped_divide(local_output, cute.make_layout(tile))
self.kernel(
tiled_buffers[0],
tiled_buffers[1],
tiled_buffers[2],
tiled_buffers[3],
tiled_buffers[4],
tiled_buffers[5],
tiled_buffers[6],
tiled_buffers[7],
tiled_input,
tiled_output,
tv_layout,
cutlass.Int32(signal),
cutlass.Int32(rank),
).launch(
grid=[num_of_elements // tile, 1, 1],
block=[tv_layout.shape[0], 1, 1],
stream=stream,
)
# GPU device kernel
@cute.kernel
def kernel(
self,
buffer0: cute.Tensor,
buffer1: cute.Tensor,
buffer2: cute.Tensor,
buffer3: cute.Tensor,
buffer4: cute.Tensor,
buffer5: cute.Tensor,
buffer6: cute.Tensor,
buffer7: cute.Tensor,
local_input: cute.Tensor,
local_output: cute.Tensor,
tv_layout: cute.Layout,
signal: cutlass.Int32,
rank: cutlass.Int32,
):
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
ping = signal % 3
pong = (signal + 1) % 3
buffers = [
buffer0,
buffer1,
buffer2,
buffer3,
buffer4,
buffer5,
buffer6,
buffer7,
]
def get_buffer():
t = buffers[2]
if rank == cutlass.Int32(0):
t = buffers[0]
if rank == cutlass.Int32(1):
t = buffers[1]
if rank == cutlass.Int32(2):
t = buffers[2]
if rank == cutlass.Int32(3):
t = buffers[3]
if rank == cutlass.Int32(4):
t = buffers[4]
if rank == cutlass.Int32(5):
t = buffers[5]
if rank == cutlass.Int32(6):
t = buffers[6]
if rank == cutlass.Int32(7):
t = buffers[7]
return t
buffer_local = get_buffer()
cta_coord = (None, bidx)
local_tile_in = local_input[cta_coord]
local_tile_out = local_output[cta_coord]
ping_coord = ((None, bidx), None, ping)
read_buffer = buffer_local[ping_coord]
pong_coord = ((None, bidx), None, pong)
clear_buffer = buffer_local[pong_coord]
write_coord = ((None, bidx), rank, ping)
write_buffers = [buffer[write_coord] for buffer in buffers]
# assume all buffers have the same element type with input
copy_atom_load = cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
buffer0.element_type,
num_bits_per_copy=64,
memory_scope=cute.nvgpu.common.MemoryScope.SYS,
memory_order=cute.nvgpu.common.MemoryOrder.VOLATILE,
)
copy_atom_store = cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
buffer0.element_type,
num_bits_per_copy=64,
memory_scope=cute.nvgpu.common.MemoryScope.SYS,
memory_order=cute.nvgpu.common.MemoryOrder.VOLATILE,
)
tiled_copy = cute.make_tiled_copy_tv(copy_atom_load, tv_layout[0], tv_layout[1])
thr_copy = tiled_copy.get_slice(tidx)
thr_write_buffer_list = [
thr_copy.partition_D(tensor) for tensor in write_buffers
]
thr_read_buffer = thr_copy.partition_S(read_buffer)
thr_clear_buffer = thr_copy.partition_D(clear_buffer)
thr_in = thr_copy.partition_S(local_tile_in)
thr_out = thr_copy.partition_D(local_tile_out)
frg_in = cute.make_fragment_like(thr_in)
frg_clear = cute.make_fragment_like(thr_clear_buffer)
frg_acc = cute.make_fragment_like(thr_out)
frg_acc.fill(0.0)
clear_tensor = frg_clear.load()
frg_size = cute.size(clear_tensor.shape)
neg0_i32_vec = cute.full_like(clear_tensor, 0x80000000, cutlass.Int32)
neg0_f32_vec = vector.bitcast(T.vector(frg_size, T.f32()), neg0_i32_vec)
neg0_f32_tensor = cute.TensorSSA(
neg0_f32_vec, clear_tensor.shape, cutlass.Float32
)
frg_clear.store(neg0_f32_tensor)
cute.copy(copy_atom_load, thr_in, frg_in)
for thr_write_buffer in thr_write_buffer_list:
cute.copy(copy_atom_store, frg_in, thr_write_buffer)
cute.copy(copy_atom_store, frg_clear, thr_clear_buffer)
frg_in_vector_neg0_i32 = cute.full_like(
frg_in, cutlass.Int32(0x80000000), cutlass.Int32
)
frg_in_size = cute.size(frg_in.shape)
for i in range(WORLD_SIZE):
read_coord = (None, 0, i)
cute.copy(copy_atom_load, thr_read_buffer[read_coord], frg_in[None, 0])
frg_vector = frg_in.load()
frg_vector_i32 = vector.bitcast(T.vector(frg_in_size, T.i32()), frg_vector)
isNotNeg0 = cute.all_(frg_vector_i32 != frg_in_vector_neg0_i32)
while not isNotNeg0:
cute.copy(copy_atom_load, thr_read_buffer[read_coord], frg_in[None, 0])
frg_vector = frg_in.load()
frg_vector_i32 = vector.bitcast(
T.vector(frg_in_size, T.i32()), frg_vector
)
isNotNeg0 = cute.all_(frg_vector_i32 != frg_in_vector_neg0_i32)
frg_acc.store(frg_in.load() + frg_acc.load())
cute.copy(copy_atom_stg, frg_acc, thr_out)
def run_all_reduce(rank, M, N, dtype: Type[cutlass.Numeric]):
setup(rank, WORLD_SIZE)
input_tensor = torch.randn(M * N, device=f"cuda:{rank}")
output_tensor = torch.zeros(M * N, device=f"cuda:{rank}")
# init tensors on different devices
t = symm_mem.empty(
[
PING_PONG_SIZE,
WORLD_SIZE,
M * N,
],
device="cuda",
).neg_()
hdl = symm_mem.rendezvous(t, dist.group.WORLD)
buffer_tensor_list = [
hdl.get_buffer(rank, t.shape, t.dtype).permute(2, 1, 0)
for rank in range(WORLD_SIZE)
]
# enable peer access
driver.cuInit(0)
dev_list = [driver.cuDeviceGet(i)[1] for i in range(WORLD_SIZE)]
ctx_list = [driver.cuDevicePrimaryCtxRetain(dev)[1] for dev in dev_list]
for i in range(WORLD_SIZE):
driver.cuCtxSetCurrent(ctx_list[i])
for j in range(WORLD_SIZE):
if i == j:
continue
driver.cuCtxEnablePeerAccess(ctx_list[j], 0)
driver.cuCtxSetCurrent(ctx_list[rank])
stream = cutlass.cuda.default_stream()
all_reduce_kernel = AllReduceKernel()
dlpack_buffers = [from_dlpack(x, assumed_align=32) for x in buffer_tensor_list]
all_reduce_kernel(
rank,
0,
from_dlpack(input_tensor, assumed_align=32),
from_dlpack(output_tensor, assumed_align=32),
*dlpack_buffers,
stream,
)
torch.cuda.synchronize(0)
# use torch api to get reference and inplace stored to input_tensor
ref_tensor = input_tensor.clone()
dist.all_reduce(ref_tensor, op=dist.ReduceOp.SUM)
# check result of output tensor, allow small error due to different accumulator datatypes
equal_mask = (ref_tensor.cpu() - output_tensor.cpu()).abs() < 1e-4
result = (equal_mask.sum()).item() == ref_tensor.numel()
if result:
print(f"rank {rank} test passed")
else:
print(f"rank {rank} test failed")
print(
"ref_tensor[ref_tensor != output_tensor]: ",
ref_tensor[ref_tensor != output_tensor],
)
print(
"output_tensor[ref_tensor != output_tensor]: ",
output_tensor[ref_tensor != output_tensor],
)
cleanup()
def main():
M = 1024
N = 1024
# each process will run run_all_reduce on different device
mp.spawn(run_all_reduce, args=(M, N, cutlass.Float32), nprocs=WORLD_SIZE, join=True)
return
if __name__ == "__main__":
main()
examples/python/CuTeDSL/ampere/call_bypass_dlpack.py
+166
View File
@@ -0,0 +1,166 @@
# 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
from typing import Tuple
import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import make_ptr
"""
An Example demonstrating how to call off-the-shelf kernel by-passing dlpack protocol
The example shows how to directly pass pointers from PyTorch tensors to off-the-shelf kernels
written by CuTe DSL with a thin customized wrapper jit function. The jit function will be
compiled with inline without introducing overhead.
To run this example:
.. code-block:: bash
python examples/ampere/call_bypass_dlpack.py
It's worth to mention that by-passing dlpack protocol can resolve the issue that dlpack doesn't handle shape-1
mode correctly. For example, the following code will fail, because dlpack will convert the shape-1 mode
with stride-1 which propagate alignment incorrectly.
.. code-block:: python
@cute.kernel
def fails_kernel(gX: cute.Tensor):
bidx, _, _ = cute.arch.block_idx()
mX = gX[None, bidx, None] # We wish to retain alignment
# assert mX.iterator.alignment == 16
@cute.jit
def fails(gX_: cute.Tensor):
gX = gX_
fails_kernel(gX).launch(grid=(1, 1, 1), block=(128, 1, 1))
gX_torch = torch.rand((128, 1, 128), device="cuda", dtype=torch.bfloat16)
fails(from_dlpack(gX_torch, assumed_align=16))
"""
# Add the current directory to sys.path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from tensorop_gemm import TensorOpGemm
@cute.jit
def tensor_op_gemm_wrapper(
a_ptr: cute.Pointer,
b_ptr: cute.Pointer,
c_ptr: cute.Pointer,
m: cutlass.Int32,
n: cutlass.Int32,
k: cutlass.Int32,
l: cutlass.Int32,
):
print(f"\n[DSL INFO] Input Parameters:")
print(f"[DSL INFO] mnkl: {(m, n, k, l)}")
# Assume alignment of shape to call tensorop_gemm example
m = cute.assume(m, divby=8)
n = cute.assume(n, divby=8)
# Torch is row major
a_layout = cute.make_ordered_layout((m, k, l), order=(0, 1, 2))
b_layout = cute.make_ordered_layout((n, k, l), order=(0, 1, 2))
c_layout = cute.make_ordered_layout((m, n, l), order=(1, 0, 2))
mA = cute.make_tensor(a_ptr, layout=a_layout)
mB = cute.make_tensor(b_ptr, layout=b_layout)
mC = cute.make_tensor(c_ptr, layout=c_layout)
print(f"[DSL INFO] mA: {mA}")
print(f"[DSL INFO] mB: {mB}")
print(f"[DSL INFO] mC: {mC}")
tensor_op_gemm = TensorOpGemm(
a_ptr.value_type, c_ptr.value_type, cutlass.Float32, (2, 2, 1)
)
print(f"\n[DSL INFO] Created TensorOpGemm instance")
print(f"[DSL INFO] Input dtype: {a_ptr.value_type}")
print(f"[DSL INFO] Output dtype: {c_ptr.value_type}")
print(f"[DSL INFO] Accumulation dtype: {cutlass.Float32}")
print(f"[DSL INFO] Atom layout: {(2, 2, 1)}")
# No need to compile inside jit function
tensor_op_gemm(mA, mB, mC)
print(f"\n[DSL INFO] Executed TensorOpGemm")
def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
print(f"\nRunning TensorOpGemm test with:")
print(f"Tensor dimensions: {mnkl}")
# (M,K,L)
a = torch.randn(
mnkl[3], mnkl[2], mnkl[0], dtype=torch.float16, device="cuda"
).permute(2, 1, 0)
# (N,K,L)
b = torch.randn(
mnkl[3], mnkl[2], mnkl[1], dtype=torch.float16, device="cuda"
).permute(2, 1, 0)
# (N,M,L)
c = torch.randn(
mnkl[3], mnkl[0], mnkl[1], dtype=torch.float16, device="cuda"
).permute(1, 2, 0)
print(f"Input tensor shapes:")
print(f"a: {a.shape}, dtype: {a.dtype}")
print(f"b: {b.shape}, dtype: {b.dtype}")
print(f"c: {c.shape}, dtype: {c.dtype}\n")
a_ptr = make_ptr(
cutlass.Float16, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
b_ptr = make_ptr(
cutlass.Float16, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
c_ptr = make_ptr(
cutlass.Float16, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
tensor_op_gemm_wrapper(a_ptr, b_ptr, c_ptr, *mnkl)
torch.cuda.synchronize()
ref = torch.einsum("mkl,nkl->mnl", a, b)
torch.testing.assert_close(c, ref, atol=1e-05, rtol=1e-05)
print(f"\n[DSL INFO] Results verified successfully!")
print(f"First few elements of result: \n{c[:3, :3, :3]}")
if __name__ == "__main__":
run_tensor_op_gemm_wrapper((512, 256, 128, 16))
examples/python/CuTeDSL/ampere/call_from_jit.py
+3 -3
View File
@@ -226,15 +226,15 @@ def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
print(f"c: {c.shape}, dtype: {c.dtype}\n")
buffer_a = BufferWithLayout(
make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem),
make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=32),
(2, 1, 0),
)
buffer_b = BufferWithLayout(
make_ptr(ab_dtype, b.data_ptr(), cute.AddressSpace.gmem),
make_ptr(ab_dtype, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=32),
(2, 1, 0),
)
buffer_c = BufferWithLayout(
make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem),
make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=32),
(2, 1, 0),
)
examples/python/CuTeDSL/ampere/distributed_vector_add.py
+189
View File
@@ -0,0 +1,189 @@
import os
import torch
import argparse
from typing import Type
from cuda.bindings import driver
import torch.distributed as dist
import torch.distributed._symmetric_memory as symm_mem
import torch.multiprocessing as mp
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
def setup(rank, world_size):
# set environment variables for torch.distributed environment
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "12995"
dist.init_process_group("nccl", rank=rank, world_size=world_size)
torch.cuda.set_device(rank)
def cleanup():
dist.destroy_process_group()
@cute.kernel
def vector_add_kernel(
g0: cute.Tensor,
g1: cute.Tensor,
g2: cute.Tensor,
g3: cute.Tensor,
g4: cute.Tensor,
g5: cute.Tensor,
g6: cute.Tensor,
g7: cute.Tensor,
gOut: cute.Tensor,
tv_layout: cute.Layout,
):
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
cta_coord = (None, bidx)
local_tile_out = gOut[cta_coord]
local_tile_list = [
g0[cta_coord],
g1[cta_coord],
g2[cta_coord],
g3[cta_coord],
g4[cta_coord],
g5[cta_coord],
g6[cta_coord],
g7[cta_coord],
]
copy_atom_load = cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
g0.element_type,
memory_order=cute.nvgpu.common.MemoryOrder.VOLATILE,
memory_scope=cute.nvgpu.common.MemoryScope.SYS,
)
copy_atom_store = cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
g0.element_type,
memory_order=cute.nvgpu.common.MemoryOrder.VOLATILE,
memory_scope=cute.nvgpu.common.MemoryScope.SYS,
)
tiled_copy = cute.make_tiled_copy_tv(copy_atom_load, tv_layout[0], tv_layout[1])
thr_copy = tiled_copy.get_slice(tidx)
thr_tensor_list = [thr_copy.partition_S(tensor) for tensor in local_tile_list]
thr_out = thr_copy.partition_D(local_tile_out)
frg_tensor_list = [cute.make_fragment_like(tensor) for tensor in thr_tensor_list]
frg_acc = cute.make_fragment_like(thr_out)
frg_acc.fill(0.0)
for thr, frg in zip(thr_tensor_list, frg_tensor_list):
cute.copy(copy_atom_load, thr, frg)
tmp = frg.load() + frg_acc.load()
frg_acc.store(tmp)
cute.copy(copy_atom_store, frg_acc, thr_out)
@cute.jit
def vector_add(
m0: cute.Tensor,
m1: cute.Tensor,
m2: cute.Tensor,
m3: cute.Tensor,
m4: cute.Tensor,
m5: cute.Tensor,
m6: cute.Tensor,
m7: cute.Tensor,
output: cute.Tensor,
):
# define constants for future use
num_of_elements = cute.size(m0.layout)
# 128 threads per block and 4 elements per thread
tv_layout = cute.make_layout(((128), (4)), stride=((1), (1)))
tile = cute.size(tv_layout.shape)
tensors = [m0, m1, m2, m3, m4, m5, m6, m7]
divided_tensors = [
cute.zipped_divide(tensor, cute.make_layout(tile)) for tensor in tensors
]
gOut = cute.zipped_divide(output, cute.make_layout(tile)) # ((Tile),(Rest))
vector_add_kernel(
divided_tensors[0],
divided_tensors[1],
divided_tensors[2],
divided_tensors[3],
divided_tensors[4],
divided_tensors[5],
divided_tensors[6],
divided_tensors[7],
gOut,
tv_layout,
).launch(
grid=[num_of_elements // tile, 1, 1],
block=[tv_layout.shape[0], 1, 1],
)
def run_vector_add(rank, world_size, M, N, dtype: Type[cutlass.Numeric]):
setup(rank, world_size)
t = symm_mem.empty(M * N, device="cuda")
hdl = symm_mem.rendezvous(t, dist.group.WORLD)
# get tensors from other devices from the symmetric memory
tensor_list = [hdl.get_buffer(rank, t.shape, t.dtype) for rank in range(world_size)]
tensor_list[rank].random_(0, 100)
# enable peer access
driver.cuInit(0)
dev_list = [driver.cuDeviceGet(i)[1] for i in range(world_size)]
ctx_list = [driver.cuDevicePrimaryCtxRetain(dev)[1] for dev in dev_list]
driver.cuCtxSetCurrent(ctx_list[rank])
for i in range(world_size):
if i == rank:
continue
driver.cuCtxEnablePeerAccess(ctx_list[i], 0)
output = torch.zeros(M * N, device=f"cuda:{rank}")
# we have to explicitly pass each tensor instead of a list of tensors
vector_add(
from_dlpack(tensor_list[0], assumed_align=32),
from_dlpack(tensor_list[1], assumed_align=32),
from_dlpack(tensor_list[2], assumed_align=32),
from_dlpack(tensor_list[3], assumed_align=32),
from_dlpack(tensor_list[4], assumed_align=32),
from_dlpack(tensor_list[5], assumed_align=32),
from_dlpack(tensor_list[6], assumed_align=32),
from_dlpack(tensor_list[7], assumed_align=32),
from_dlpack(output, assumed_align=32),
)
sum_tensor = sum([tensor.cpu() for tensor in tensor_list])
if sum(sum_tensor.cpu() == output.cpu()) == sum_tensor.numel():
print("test passed")
else:
print("test failed")
print(sum_tensor.cpu())
print(output.cpu())
cleanup()
def main():
world_size = torch.cuda.device_count()
M = 1024
N = 1024
# each process will run run_vector_add on different device
mp.spawn(
run_vector_add,
args=(world_size, M, N, cutlass.Float32),
nprocs=world_size,
join=True,
)
return
if __name__ == "__main__":
main()
examples/python/CuTeDSL/ampere/dynamic_smem_size.py
+114
View File
@@ -0,0 +1,114 @@
# 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 cutlass
"""
Example of automatic shared memory size computation for configuring kernel launch
This example demonstrates how to let the DSL automatically set shared memory
size for a kernel launch rather explicitly configuring it at launch time,
provided that developers are using `SmemAllocator` for all allocations.
Usage:
python dynamic_smem_size.py # Show auto inference
"""
@cute.struct
class SharedData:
"""A struct to demonstrate shared memory allocation."""
values: cute.struct.MemRange[cutlass.Float32, 64] # 256 bytes
counter: cutlass.Int32 # 4 bytes
flag: cutlass.Int8 # 1 byte
@cute.kernel
def kernel():
"""
Example kernel that allocates shared memory.
The total allocation will be automatically calculated when smem=None.
"""
allocator = cutlass.utils.SmemAllocator()
# Allocate various types of shared memory
shared_data = allocator.allocate(SharedData)
raw_buffer = allocator.allocate(512, byte_alignment=64)
int_array = allocator.allocate_array(element_type=cutlass.Int32, num_elems=128)
tensor_smem = allocator.allocate_tensor(
element_type=cutlass.Float16,
layout=cute.make_layout((32, 16)),
byte_alignment=16,
swizzle=None,
)
return
@cute.kernel
def kernel_no_smem():
"""
Example kernel that does not allocates shared memory.
The total allocation will be automatically calculated as 0 when smem=None.
"""
tidx, _, _ = cute.arch.block_idx()
if tidx == 0:
cute.printf("Hello world")
return
if __name__ == "__main__":
# Initialize CUDA context
cutlass.cuda.initialize_cuda_context()
print("Launching kernel with auto smem size. (launch config `smem=None`)")
# Compile the example
@cute.jit
def launch_kernel1():
k = kernel()
k.launch(
grid=(1, 1, 1),
block=(1, 1, 1),
)
print(f"Kernel recorded internal smem usage: {k.smem_usage()}")
@cute.jit
def launch_kernel2():
k = kernel_no_smem()
k.launch(
grid=(1, 1, 1),
block=(1, 1, 1),
)
print(f"Kernel recorded internal smem usage: {k.smem_usage()}")
cute.compile(launch_kernel1)
cute.compile(launch_kernel2)
print("PASS")
examples/python/CuTeDSL/ampere/flash_attention_v2.py
+7 -30
View File
@@ -327,7 +327,6 @@ class FlashAttentionForwardAmpere:
).launch(
grid=grid_dim,
block=[self._num_threads, 1, 1],
smem=SharedStorage.size_in_bytes(),
stream=stream,
)
@@ -1014,13 +1013,10 @@ class FlashAttentionForwardAmpere:
)
# compute exp(x - max) using exp2(x * log_2(e) - max * log_2(e))
acc_S_row_exp = cute.TensorSSA(
self._exp2f(
acc_S_row * softmax_params.softmax_scale_log2
- row_max_cur_row * softmax_params.softmax_scale_log2
),
tuple(acc_S_row.shape),
cutlass.Float32,
acc_S_row_exp = cute.math.exp2(
acc_S_row * softmax_params.softmax_scale_log2
- row_max_cur_row * softmax_params.softmax_scale_log2,
fastmath=True,
)
# acc_S_row_sum => f32
acc_S_row_sum = acc_S_row_exp.reduce(
@@ -1028,9 +1024,10 @@ class FlashAttentionForwardAmpere:
)
# if it is not the first tile, load the row r of previous row_max and minus row_max_cur_row to update row_sum.
if cutlass.const_expr(not is_first_n_block):
prev_minus_cur_exp = self._exp2f(
prev_minus_cur_exp = cute.math.exp2(
row_max_prev_row * softmax_params.softmax_scale_log2
- row_max_cur_row * softmax_params.softmax_scale_log2
- row_max_cur_row * softmax_params.softmax_scale_log2,
fastmath=True,
)
acc_S_row_sum = (
acc_S_row_sum + softmax_params.row_sum[r] * prev_minus_cur_exp
@@ -1141,26 +1138,6 @@ class FlashAttentionForwardAmpere:
"""
return self._threadquad_reduce(val, lambda x, y: x + y)
def _exp2f(
self, x: Union[cute.TensorSSA, cutlass.Float32]
) -> Union[cute.TensorSSA, cutlass.Float32]:
"""exp2f calculation for both vector and scalar.
:param x: input value
:type x: cute.TensorSSA or cutlass.Float32
:return: exp2 value
:rtype: cute.TensorSSA or cutlass.Float32
"""
if isinstance(x, cute.TensorSSA):
res = cute.make_fragment(x.shape, cutlass.Float32)
res.store(x)
for i in range(cute.size(x.shape)):
res[i] = self._exp2f(res[i])
return res.load()
return cute.arch.exp2(x)
def run(
dtype: Type[cutlass.Numeric],
examples/python/CuTeDSL/ampere/sgemm.py
+10 -8
View File
@@ -136,10 +136,6 @@ class SGemm:
stride=(1, (self._bN + padding_b), self._bK * (self._bN + padding_b)),
)
smem_size = cute.size_in_bytes(mA.element_type, sA_layout) + cute.size_in_bytes(
mB.element_type, sB_layout
)
# ///////////////////////////////////////////////////////////////////////////////
# Create copy layouts that will be used for asynchronous
# global memory -> shared memory copies:
@@ -258,7 +254,6 @@ class SGemm:
).launch(
grid=grid_dim,
block=[cute.size(atoms_layout), 1, 1],
smem=smem_size,
stream=stream,
)
@@ -738,14 +733,20 @@ def run(
print("Compiling kernel with cute.compile ...")
start_time = time.time()
gemm = cute.compile(sgemm, a_tensor, b_tensor, c_tensor, stream=current_stream)
compiled_fn = cute.compile(
sgemm,
a_tensor,
b_tensor,
c_tensor,
stream=current_stream,
)
compilation_time = time.time() - start_time
print(f"Compilation time: {compilation_time:.4f} seconds")
print("Executing GEMM kernel...")
if not skip_ref_check:
gemm(a_tensor, b_tensor, c_tensor)
compiled_fn(a_tensor, b_tensor, c_tensor)
torch.cuda.synchronize()
print("Verifying results...")
ref = torch.einsum("mk,nk->mn", a, b)
@@ -804,7 +805,7 @@ def run(
)
avg_time_us = testing.benchmark(
gemm,
compiled_fn,
workspace_generator=generate_tensors,
workspace_count=workspace_count,
stream=current_stream,
@@ -837,6 +838,7 @@ if __name__ == "__main__":
parser.add_argument("--c_major", choices=["n", "m"], default="n")
parser.add_argument("--warmup_iterations", default=2, type=int)
parser.add_argument("--iterations", default=100, type=int)
parser.add_argument("--static_shape", action="store_true")
parser.add_argument("--skip_ref_check", action="store_true")
parser.add_argument(
"--use_cold_l2",
examples/python/CuTeDSL/ampere/smem_allocator.py
+1 -1
View File
@@ -69,7 +69,7 @@ class complex:
class SharedStorage:
# struct elements with natural alignment
a: cute.struct.MemRange[cutlass.Float32, 32] # array
b: cutlass.Int64 # scalar
b: cutlass.Int64 # saclar
c: complex # nested struct
# struct elements with strict alignment
x: cute.struct.Align[
examples/python/CuTeDSL/ampere/tensorop_gemm.py
+2 -2
View File
@@ -471,7 +471,7 @@ class TensorOpGemm:
cute.arch.sync_threads()
# Start async loads for the first k-tile. Here we take care of the k residue
# via if/else check along the k dimension. Because we shifted the identity tensor
# by the residue_k and because the identity tensor is a counting tensor, the
# by the residue_k and because the identity tensor is a coord tensor, the
# values of any identity tensor element that is poison is less than -1
num_smem_stages = cute.size(tAsA, mode=[3])
k_tile_count = cute.size(tAgA, mode=[3])
@@ -683,7 +683,7 @@ class TensorOpGemm:
# Copy results of D back to shared memory
cute.autovec_copy(tCrD, tCsC)
# Create counting tensor for C
# Create coord tensor for C
ceilM, ceilN, _ = cute.ceil_div(mC.shape, (self.bM, self.bN, 1))
mcC = cute.make_identity_tensor(
(
examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py
+21 -22
View File
@@ -610,7 +610,6 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
smem=self.shared_storage.size_in_bytes(),
stream=stream,
)
return
@@ -797,7 +796,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
gC_mnl = cute.local_tile(
mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
)
k_block_cnt = cute.size(gA_mkl, mode=[3])
k_tile_cnt = cute.size(gA_mkl, mode=[3])
#
# Partition global tensor for TiledMMA_A/B/C
@@ -946,17 +945,17 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])
]
# Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt
# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt
ab_producer_state.reset_count()
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < k_block_cnt:
if ab_producer_state.count < k_tile_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(
ab_producer_state
)
#
# Tma load loop
#
for k_block in cutlass.range(0, k_block_cnt, 1, unroll=1):
for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
# Conditionally wait for AB buffer empty
ab_pipeline.producer_acquire(
ab_producer_state, peek_ab_empty_status
@@ -992,10 +991,10 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
mcast_mask=sfb_full_mcast_mask,
)
# Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt + k_block + 1
# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1
ab_producer_state.advance()
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < k_block_cnt:
if ab_producer_state.count < k_tile_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(
ab_producer_state
)
@@ -1103,10 +1102,10 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
# (MMA, MMA_M, MMA_N)
tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)]
# Peek (try_wait) AB buffer full for k_block = 0
# Peek (try_wait) AB buffer full for k_tile = 0
ab_consumer_state.reset_count()
peek_ab_full_status = cutlass.Boolean(1)
if ab_consumer_state.count < k_block_cnt and is_leader_cta:
if ab_consumer_state.count < k_tile_cnt and is_leader_cta:
peek_ab_full_status = ab_pipeline.consumer_try_wait(
ab_consumer_state
)
@@ -1125,7 +1124,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
#
# Mma mainloop
#
for k_block in range(k_block_cnt):
for k_tile in range(k_tile_cnt):
if is_leader_cta:
# Conditionally wait for AB buffer full
ab_pipeline.consumer_wait(
@@ -1154,44 +1153,44 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
)
# tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB
num_kphases = cute.size(tCrA, mode=[2])
for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
kphase_coord = (
num_kblocks = cute.size(tCrA, mode=[2])
for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
kblock_coord = (
None,
None,
kphase_idx,
kblock_idx,
ab_consumer_state.index,
)
# Set SFA/SFB tensor to tiled_mma
sf_kphase_coord = (None, None, kphase_idx)
sf_kblock_coord = (None, None, kblock_idx)
tiled_mma.set(
tcgen05.Field.SFA,
tCtSFA[sf_kphase_coord].iterator,
tCtSFA[sf_kblock_coord].iterator,
)
tiled_mma.set(
tcgen05.Field.SFB,
tCtSFB[sf_kphase_coord].iterator,
tCtSFB[sf_kblock_coord].iterator,
)
cute.gemm(
tiled_mma,
tCtAcc,
tCrA[kphase_coord],
tCrB[kphase_coord],
tCrA[kblock_coord],
tCrB[kblock_coord],
tCtAcc,
)
# Enable accumulate on tCtAcc after first kphase
# Enable accumulate on tCtAcc after first kblock
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
# Async arrive AB buffer empty
ab_pipeline.consumer_release(ab_consumer_state)
# Peek (try_wait) AB buffer full for k_block = k_block + 1
# Peek (try_wait) AB buffer full for k_tile = k_tile + 1
ab_consumer_state.advance()
peek_ab_full_status = cutlass.Boolean(1)
if ab_consumer_state.count < k_block_cnt:
if ab_consumer_state.count < k_tile_cnt:
if is_leader_cta:
peek_ab_full_status = ab_pipeline.consumer_try_wait(
ab_consumer_state
examples/python/CuTeDSL/blackwell/dense_gemm.py
+23 -24
View File
@@ -486,7 +486,6 @@ class DenseGemmKernel:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
smem=self.shared_storage.size_in_bytes(),
stream=stream,
)
return
@@ -660,7 +659,7 @@ class DenseGemmKernel:
gC_mnl = cute.local_tile(
mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
)
k_block_cnt = cute.size(gA_mkl, mode=[3])
k_tile_cnt = cute.size(gA_mkl, mode=[3])
#
# Partition global tensor for TiledMMA_A/B/C
@@ -788,19 +787,19 @@ class DenseGemmKernel:
#
# Pipelining TMA load A/B and MMA mainloop
#
prefetch_k_block_cnt = cutlass.min(self.num_ab_stage - 2, k_block_cnt)
prefetch_k_tile_cnt = cutlass.min(self.num_ab_stage - 2, k_tile_cnt)
if warp_idx == 0:
# Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt
# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < k_block_cnt:
if ab_producer_state.count < k_tile_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(
ab_producer_state
)
#
# Prefetch TMA load A/B
#
for prefetch_idx in cutlass.range(prefetch_k_block_cnt, unroll=1):
for prefetch_idx in cutlass.range(prefetch_k_tile_cnt, unroll=1):
# Conditionally wait for AB buffer empty
ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status)
@@ -820,27 +819,27 @@ class DenseGemmKernel:
mcast_mask=b_full_mcast_mask,
)
# Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt + k_block + 1
# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1
ab_producer_state.advance()
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < k_block_cnt:
if ab_producer_state.count < k_tile_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(
ab_producer_state
)
# Peek (try_wait) AB buffer full for k_block = 0
# Peek (try_wait) AB buffer full for k_tile = 0
peek_ab_full_status = cutlass.Boolean(1)
if ab_consumer_state.count < k_block_cnt and is_leader_cta:
if ab_consumer_state.count < k_tile_cnt and is_leader_cta:
peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state)
#
# MMA mainloop
#
for k_block in range(k_block_cnt):
for k_tile in range(k_tile_cnt):
# Conditionally wait for AB buffer empty
ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status)
if ab_producer_state.count < k_block_cnt:
if ab_producer_state.count < k_tile_cnt:
# TMA load A/B
cute.copy(
tma_atom_a,
@@ -862,35 +861,35 @@ class DenseGemmKernel:
ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status)
# tCtAcc += tCrA * tCrB
num_kphases = cute.size(tCrA, mode=[2])
for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
kphase_coord = (None, None, kphase_idx, ab_consumer_state.index)
num_kblocks = cute.size(tCrA, mode=[2])
for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
kblock_coord = (None, None, kblock_idx, ab_consumer_state.index)
cute.gemm(
tiled_mma,
tCtAcc,
tCrA[kphase_coord],
tCrB[kphase_coord],
tCrA[kblock_coord],
tCrB[kblock_coord],
tCtAcc,
)
# Enable accumulate on tCtAcc after first kphase
# Enable accumulate on tCtAcc after first kblock
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
# Async arrive AB buffer empty
ab_pipeline.consumer_release(ab_consumer_state)
# Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt + k_block + 1
# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1
ab_producer_state.advance()
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < k_block_cnt:
if ab_producer_state.count < k_tile_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(
ab_producer_state
)
# Peek (try_wait) AB buffer full for k_block = k_block + 1
# Peek (try_wait) AB buffer full for k_tile = k_tile + 1
ab_consumer_state.advance()
peek_ab_full_status = cutlass.Boolean(1)
if ab_consumer_state.count < k_block_cnt:
if ab_consumer_state.count < k_tile_cnt:
if is_leader_cta:
peek_ab_full_status = ab_pipeline.consumer_try_wait(
ab_consumer_state
@@ -1009,8 +1008,8 @@ class DenseGemmKernel:
# Wait A/B buffer empty
#
if warp_idx == 0:
# Reverse prefetch_k_block_cnt times to next available buffer
for i in range(prefetch_k_block_cnt):
# Reverse prefetch_k_tile_cnt times to next available buffer
for i in range(prefetch_k_tile_cnt):
ab_producer_state.reverse()
ab_pipeline.producer_tail(ab_producer_state)
return
examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py
+18 -19
View File
@@ -510,7 +510,6 @@ class PersistentDenseGemmKernel:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
smem=self.shared_storage.size_in_bytes(),
stream=stream,
)
return
@@ -669,7 +668,7 @@ class PersistentDenseGemmKernel:
gC_mnl = cute.local_tile(
mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
)
k_block_cnt = cute.size(gA_mkl, mode=[3])
k_tile_cnt = cute.size(gA_mkl, mode=[3])
#
# Partition global tensor for TiledMMA_A/B/C
@@ -774,17 +773,17 @@ class PersistentDenseGemmKernel:
(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])
]
# Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt
# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt
ab_producer_state.reset_count()
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < k_block_cnt:
if ab_producer_state.count < k_tile_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(
ab_producer_state
)
#
# Tma load loop
#
for k_block in cutlass.range(0, k_block_cnt, 1, unroll=1):
for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
# Conditionally wait for AB buffer empty
ab_pipeline.producer_acquire(
ab_producer_state, peek_ab_empty_status
@@ -806,10 +805,10 @@ class PersistentDenseGemmKernel:
mcast_mask=b_full_mcast_mask,
)
# Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt + k_block + 1
# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1
ab_producer_state.advance()
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < k_block_cnt:
if ab_producer_state.count < k_tile_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(
ab_producer_state
)
@@ -877,10 +876,10 @@ class PersistentDenseGemmKernel:
# (MMA, MMA_M, MMA_N)
tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)]
# Peek (try_wait) AB buffer full for k_block = 0
# Peek (try_wait) AB buffer full for k_tile = 0
ab_consumer_state.reset_count()
peek_ab_full_status = cutlass.Boolean(1)
if ab_consumer_state.count < k_block_cnt and is_leader_cta:
if ab_consumer_state.count < k_tile_cnt and is_leader_cta:
peek_ab_full_status = ab_pipeline.consumer_try_wait(
ab_consumer_state
)
@@ -899,7 +898,7 @@ class PersistentDenseGemmKernel:
#
# Mma mainloop
#
for k_block in range(k_block_cnt):
for k_tile in range(k_tile_cnt):
if is_leader_cta:
# Conditionally wait for AB buffer full
ab_pipeline.consumer_wait(
@@ -907,32 +906,32 @@ class PersistentDenseGemmKernel:
)
# tCtAcc += tCrA * tCrB
num_kphases = cute.size(tCrA, mode=[2])
for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
kphase_coord = (
num_kblocks = cute.size(tCrA, mode=[2])
for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
kblock_coord = (
None,
None,
kphase_idx,
kblock_idx,
ab_consumer_state.index,
)
cute.gemm(
tiled_mma,
tCtAcc,
tCrA[kphase_coord],
tCrB[kphase_coord],
tCrA[kblock_coord],
tCrB[kblock_coord],
tCtAcc,
)
# Enable accumulate on tCtAcc after first kphase
# Enable accumulate on tCtAcc after first kblock
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
# Async arrive AB buffer empty
ab_pipeline.consumer_release(ab_consumer_state)
# Peek (try_wait) AB buffer full for k_block = k_block + 1
# Peek (try_wait) AB buffer full for k_tile = k_tile + 1
ab_consumer_state.advance()
peek_ab_full_status = cutlass.Boolean(1)
if ab_consumer_state.count < k_block_cnt:
if ab_consumer_state.count < k_tile_cnt:
if is_leader_cta:
peek_ab_full_status = ab_pipeline.consumer_try_wait(
ab_consumer_state
examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py
+28 -45
View File
@@ -110,17 +110,6 @@ Constraints:
"""
class PipelineStateMinimal:
"""
Pipeline state contains an index and phase bit corresponding to the current position in the circular buffer.
"""
def __init__(self, count, index, phase):
self.count = count
self.index = index
self.phase = phase
class DenseGemmKernel:
"""
This class implements batched matrix multiplication (C = A x B) with support for various data types
@@ -497,7 +486,6 @@ class DenseGemmKernel:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
smem=self.shared_storage.size_in_bytes(),
stream=stream,
)
return
@@ -576,13 +564,19 @@ class DenseGemmKernel:
pipeline.Agent.Thread, num_tma_producer
)
ab_pipeline = pipeline.PipelineTmaUmma.create(
barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
num_stages=self.num_ab_stage,
producer_group=ab_pipeline_producer_group,
consumer_group=ab_pipeline_consumer_group,
tx_count=self.num_tma_load_bytes,
barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
cta_layout_vmnk=cluster_layout_vmnk,
)
ab_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_ab_stage
)
ab_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_ab_stage
)
# Initialize acc_pipeline (barrier) and states
acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
@@ -590,10 +584,10 @@ class DenseGemmKernel:
pipeline.Agent.Thread, self.threads_per_cta, self.threads_per_cta
)
acc_pipeline = pipeline.PipelineUmmaAsync.create(
barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
num_stages=self.num_acc_stage,
producer_group=acc_pipeline_producer_group,
consumer_group=acc_pipeline_consumer_group,
barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
cta_layout_vmnk=cluster_layout_vmnk,
)
acc_producer_state = pipeline.make_pipeline_state(
@@ -665,7 +659,7 @@ class DenseGemmKernel:
gC_mnl = cute.local_tile(
mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
)
k_block_cnt = cute.size(gA_mkl, mode=[3])
k_tile_cnt = cute.size(gA_mkl, mode=[3])
#
# Partition global tensor for TiledMMA_A/B/C
@@ -793,24 +787,12 @@ class DenseGemmKernel:
# ///////////////////////////////////////////////////////////////////////////////
# MAINLOOP
# ///////////////////////////////////////////////////////////////////////////////
prefetch_k_block_cnt = cutlass.min(self.num_ab_stage - 2, k_block_cnt)
prefetch_k_tile_cnt = cutlass.min(self.num_ab_stage - 2, k_tile_cnt)
if warp_idx == 0:
for k_block in cutlass.range(
k_block_cnt,
pipelining=self.num_ab_stage - 2,
for k_tile in cutlass.range(
k_tile_cnt,
prefetch_stages=self.num_ab_stage - 2,
):
ab_producer_state = PipelineStateMinimal(
k_block,
k_block % self.num_ab_stage,
cutlass.Int32((k_block // self.num_ab_stage) % 2) ^ 1,
)
ab_consumer_state = PipelineStateMinimal(
k_block,
k_block % self.num_ab_stage,
cutlass.Int32((k_block // self.num_ab_stage) % 2),
)
# wait for AB buffer empty
ab_pipeline.producer_acquire(ab_producer_state)
@@ -835,22 +817,26 @@ class DenseGemmKernel:
ab_pipeline.consumer_wait(ab_consumer_state)
# tCtAcc += tCrA * tCrB
num_kphases = cute.size(tCrA, mode=[2])
for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
kphase_coord = (None, None, kphase_idx, ab_consumer_state.index)
num_kblocks = cute.size(tCrA, mode=[2])
for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
kblock_coord = (None, None, kblock_idx, ab_consumer_state.index)
cute.gemm(
tiled_mma,
tCtAcc,
tCrA[kphase_coord],
tCrB[kphase_coord],
tCrA[kblock_coord],
tCrB[kblock_coord],
tCtAcc,
)
# Enable accumulate on tCtAcc after first kphase
# Enable accumulate on tCtAcc after first kblock
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
# Async arrive AB buffer empty
ab_pipeline.consumer_release(ab_consumer_state)
ab_producer_state.advance()
ab_consumer_state.advance()
# Async arrive accumulator buffer full
if is_leader_cta:
acc_pipeline.producer_commit(acc_producer_state)
@@ -964,12 +950,10 @@ class DenseGemmKernel:
# Wait A/B buffer empty
#
if warp_idx == 0:
ab_producer_state = PipelineStateMinimal(
k_block_cnt,
k_block_cnt % self.num_ab_stage,
cutlass.Int32((k_block_cnt // self.num_ab_stage) % 2) ^ 1,
)
ab_pipeline.producer_acquire(ab_producer_state)
# Reverse prefetch_k_tile_cnt times to next available buffer
for i in range(prefetch_k_tile_cnt):
ab_producer_state.reverse()
ab_pipeline.producer_tail(ab_producer_state)
return
def epilog_tmem_copy_and_partition(
@@ -1579,7 +1563,6 @@ def run_dense_gemm(
warmup_iterations: int = 0,
iterations: int = 1,
skip_ref_check: bool = False,
measure_launch_overhead=False,
):
"""
Prepare A/B/C tensors, launch GPU kernel, and reference checking.
@@ -1725,7 +1708,7 @@ def run_dense_gemm(
ref_c = ref
elif c_dtype in {cutlass.Float8E5M2, cutlass.Float8E4M3FN}:
# m major: (l, n, m) -> (m, n, l)
# k major: (l, m, n) -> (m, n, l)
# n major: (l, m, n) -> (m, n, l)
permute_order = (1, 2, 0) if c_major == "n" else (2, 1, 0)
shape = (l, m, n) if c_major == "n" else (l, n, m)
f8_torch_tensor = cutlass_torch.create_and_permute_torch_tensor(
examples/python/CuTeDSL/blackwell/fmha.py
+320 -571
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/grouped_gemm.py
+43 -44
View File
@@ -475,7 +475,6 @@ class GroupedGemmKernel:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
smem=self.shared_storage.size_in_bytes(),
stream=stream,
)
return
@@ -785,7 +784,7 @@ class GroupedGemmKernel:
)
tensormap_init_done = cutlass.Boolean(False)
# tile count we have searched
total_k_block_cnt = cutlass.Int32(0)
total_k_tile_cnt = cutlass.Int32(0)
# group index of last tile
last_group_idx = cutlass.Int32(-1)
work_tile = tile_sched.initial_work_tile_info()
@@ -795,7 +794,7 @@ class GroupedGemmKernel:
cur_tile_coord,
problem_sizes_mnkl,
)
cur_k_block_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k
cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k
cur_group_idx = grouped_gemm_cta_tile_info.group_idx
is_group_changed = cur_group_idx != last_group_idx
# skip tensormap update if we're working on the same group
@@ -861,17 +860,17 @@ class GroupedGemmKernel:
(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])
]
num_prev_k_blk = total_k_block_cnt
total_k_block_cnt += cur_k_block_cnt
num_prev_k_blk = total_k_tile_cnt
total_k_tile_cnt += cur_k_tile_cnt
# Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt
tma_wr_k_block = cutlass.Int32(0)
smem_wr_buffer = (num_prev_k_blk + tma_wr_k_block) % self.num_ab_stage
# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt
tma_wr_k_tile = cutlass.Int32(0)
smem_wr_buffer = (num_prev_k_blk + tma_wr_k_tile) % self.num_ab_stage
tma_wr_ab_empty_phase = (
num_prev_k_blk + tma_wr_k_block
num_prev_k_blk + tma_wr_k_tile
) // self.num_ab_stage % 2 ^ 1
peek_ab_empty_status = cute.arch.mbarrier_conditional_try_wait(
tma_wr_k_block < cur_k_block_cnt,
tma_wr_k_tile < cur_k_tile_cnt,
ab_empty_mbar_ptr + smem_wr_buffer,
tma_wr_ab_empty_phase,
)
@@ -882,10 +881,10 @@ class GroupedGemmKernel:
#
# Tma load loop
#
for k_block in cutlass.range(0, cur_k_block_cnt, 1, unroll=1):
tma_wr_k_block_next = tma_wr_k_block + 1
for k_tile in cutlass.range(0, cur_k_tile_cnt, 1, unroll=1):
tma_wr_k_tile_next = tma_wr_k_tile + 1
smem_wr_buffer_next = (
num_prev_k_blk + tma_wr_k_block_next
num_prev_k_blk + tma_wr_k_tile_next
) % self.num_ab_stage
tma_wr_ab_empty_phase_next = (
tma_wr_ab_empty_phase ^ 1
@@ -911,7 +910,7 @@ class GroupedGemmKernel:
# Load A/B with TMA
cute.copy(
tma_atom_a,
tAgA_slice[(None, tma_wr_k_block)],
tAgA_slice[(None, tma_wr_k_tile)],
tAsA[(None, smem_wr_buffer)],
tma_bar_ptr=smem_full_mbar_ptr,
mcast_mask=a_full_mcast_mask,
@@ -922,7 +921,7 @@ class GroupedGemmKernel:
)
cute.copy(
tma_atom_b,
tBgB_slice[(None, tma_wr_k_block)],
tBgB_slice[(None, tma_wr_k_tile)],
tBsB[(None, smem_wr_buffer)],
tma_bar_ptr=smem_full_mbar_ptr,
mcast_mask=b_full_mcast_mask,
@@ -932,14 +931,14 @@ class GroupedGemmKernel:
),
)
# Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt + k_block + 1
# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1
peek_ab_empty_status = cute.arch.mbarrier_conditional_try_wait(
tma_wr_k_block_next < cur_k_block_cnt,
tma_wr_k_tile_next < cur_k_tile_cnt,
ab_empty_mbar_ptr + smem_wr_buffer_next,
tma_wr_ab_empty_phase_next,
)
tma_wr_k_block = tma_wr_k_block_next
tma_wr_k_tile = tma_wr_k_tile_next
smem_wr_buffer = smem_wr_buffer_next
tma_wr_ab_empty_phase = tma_wr_ab_empty_phase_next
@@ -998,12 +997,12 @@ class GroupedGemmKernel:
work_tile = tile_sched.initial_work_tile_info()
# tile count we have searched
total_k_block_cnt = cutlass.Int32(0)
total_k_tile_cnt = cutlass.Int32(0)
while work_tile.is_valid_tile:
cur_tile_coord = work_tile.tile_idx
# MMA warp is only interested in number of tiles along K dimension
(
cur_k_block_cnt,
cur_k_tile_cnt,
cur_group_idx,
) = group_gemm_ts_helper.search_cluster_tile_count_k(
cur_tile_coord,
@@ -1014,17 +1013,17 @@ class GroupedGemmKernel:
# (MMA, MMA_M, MMA_N)
tCtAcc = tCtAcc_base[(None, None, None, acc_buf_idx)]
num_prev_k_blk = total_k_block_cnt
total_k_block_cnt += cur_k_block_cnt
num_prev_k_blk = total_k_tile_cnt
total_k_tile_cnt += cur_k_tile_cnt
# Peek (try_wait) AB buffer full for k_block = 0
mma_rd_k_block = cutlass.Int32(0)
smem_rd_buffer = (num_prev_k_blk + mma_rd_k_block) % self.num_ab_stage
# Peek (try_wait) AB buffer full for k_tile = 0
mma_rd_k_tile = cutlass.Int32(0)
smem_rd_buffer = (num_prev_k_blk + mma_rd_k_tile) % self.num_ab_stage
need_check_rd_buffer_full = (
mma_rd_k_block < cur_k_block_cnt and is_leader_cta
mma_rd_k_tile < cur_k_tile_cnt and is_leader_cta
)
mma_rd_ab_full_phase = (
(num_prev_k_blk + mma_rd_k_block) // self.num_ab_stage % 2
(num_prev_k_blk + mma_rd_k_tile) // self.num_ab_stage % 2
)
peek_ab_full_status = cute.arch.mbarrier_conditional_try_wait(
need_check_rd_buffer_full,
@@ -1051,10 +1050,10 @@ class GroupedGemmKernel:
#
# Mma mainloop
#
for k_block in range(cur_k_block_cnt):
mma_rd_k_block_next = cutlass.Int32(k_block + 1)
for k_tile in range(cur_k_tile_cnt):
mma_rd_k_tile_next = cutlass.Int32(k_tile + 1)
smem_rd_buffer_next = (
num_prev_k_blk + mma_rd_k_block_next
num_prev_k_blk + mma_rd_k_tile_next
) % self.num_ab_stage
mma_rd_ab_full_phase_next = (
mma_rd_ab_full_phase ^ 1
@@ -1069,18 +1068,18 @@ class GroupedGemmKernel:
)
# tCtAcc += tCrA * tCrB
num_kphases = cute.size(tCrA, mode=[2])
for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
kphase_coord = (None, None, kphase_idx, smem_rd_buffer)
num_kblocks = cute.size(tCrA, mode=[2])
for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
kblock_coord = (None, None, kblock_idx, smem_rd_buffer)
cute.gemm(
tiled_mma,
tCtAcc,
tCrA[kphase_coord],
tCrB[kphase_coord],
tCrA[kblock_coord],
tCrB[kblock_coord],
tCtAcc,
)
# Enable accumulate on tCtAcc after first kphase
# Enable accumulate on tCtAcc after first kblock
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
# Async arrive AB buffer empty
@@ -1091,9 +1090,9 @@ class GroupedGemmKernel:
self.cta_group,
)
# Peek (try_wait) AB buffer full for k_block = k_block + 1
# Peek (try_wait) AB buffer full for k_tile = k_tile + 1
need_check_rd_buffer_full = (
mma_rd_k_block_next < cur_k_block_cnt and is_leader_cta
mma_rd_k_tile_next < cur_k_tile_cnt and is_leader_cta
)
peek_ab_full_status = cute.arch.mbarrier_conditional_try_wait(
@@ -1102,7 +1101,7 @@ class GroupedGemmKernel:
mma_rd_ab_full_phase_next,
)
mma_rd_k_block = mma_rd_k_block_next
mma_rd_k_tile = mma_rd_k_tile_next
smem_rd_buffer = smem_rd_buffer_next
mma_rd_ab_full_phase = mma_rd_ab_full_phase_next
@@ -1201,7 +1200,7 @@ class GroupedGemmKernel:
# wait tensormap initialization complete before update
tensormap_manager.fence_tensormap_initialization()
# tile count we have searched
total_k_block_cnt = cutlass.Int32(0)
total_k_tile_cnt = cutlass.Int32(0)
# group index of last tile
last_group_idx = cutlass.Int32(-1)
while work_tile.is_valid_tile:
@@ -1240,8 +1239,8 @@ class GroupedGemmKernel:
grouped_gemm_cta_tile_info.cta_tile_idx_n,
0,
)
cur_k_block_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k
total_k_block_cnt += cur_k_block_cnt
cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k
total_k_tile_cnt += cur_k_tile_cnt
#
# Slice to per mma tile index
@@ -1370,8 +1369,8 @@ class GroupedGemmKernel:
#
if warp_idx == self.epilog_warp_id[0]:
cute.arch.mbarrier_wait(
(ab_empty_mbar_ptr + ((total_k_block_cnt - 1) % self.num_ab_stage)),
(((total_k_block_cnt - 1) // self.num_ab_stage) % 2),
(ab_empty_mbar_ptr + ((total_k_tile_cnt - 1) % self.num_ab_stage)),
(((total_k_tile_cnt - 1) // self.num_ab_stage) % 2),
)
@cute.jit
examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py
+18 -12
View File
@@ -622,7 +622,6 @@ class SSDKernel:
block=[self.threads_per_cta, 1, 1],
cluster=self.cluster_shape_mnk,
min_blocks_per_mp=1,
smem=self.shared_storage.size_in_bytes(),
stream=stream,
)
@@ -693,7 +692,7 @@ class SSDKernel:
G = cute.size(tma_tensor_b, mode=[3])
NGROUP_RATIO = EH // G
# Make tiledMma
# Make TiledMma
(
tiled_mma_intra1,
tiled_mma_intra2,
@@ -1745,7 +1744,7 @@ class SSDKernel:
cute.arch.fence_view_async_tmem_load()
# Combine INTER1_ACC/last_column/State
exp_last_column = cute.arch.exp(last_column.ir_value())
exp_last_column = cute.math.exp(last_column, fastmath=True)
for reg_idx in range(0, cute.size(tTR_rP), 2):
(
tTR_rP[reg_idx],
@@ -2267,9 +2266,11 @@ class SSDKernel:
) = cute.arch.fma_packed_f32x2(
(tTR_rInter[reg_idx], tTR_rInter[reg_idx + 1]),
(
cute.arch.exp(tTR_rDeltaA[reg_idx].ir_value()),
cute.arch.exp(
tTR_rDeltaA[reg_idx + 1].ir_value()
cute.math.exp(
tTR_rDeltaA[reg_idx], fastmath=True
),
cute.math.exp(
tTR_rDeltaA[reg_idx + 1], fastmath=True
),
),
(tTR_rIntra[reg_idx], tTR_rIntra[reg_idx + 1]),
@@ -3072,14 +3073,19 @@ class SSDKernel:
m, n = tCoord[subtile_idx]
if m < n:
tCompute[subtile_idx] = cutlass.Float32(-float("inf"))
LOG2_E = cutlass.Float32(1.4426950408889634)
for subtile_idx in cutlass.range(0, cute.size(tTR_rQ), 2, unroll_full=True):
# TODO: use math.exp directly
tCompute_log2e = cute.arch.mul_packed_f32x2(
(tCompute[subtile_idx], tCompute[subtile_idx + 1]), (LOG2_E, LOG2_E)
)
(
tCompute[subtile_idx],
tCompute[subtile_idx + 1],
) = cute.arch.mul_packed_f32x2(
cute.arch.exp_packed_f32x2(
(tCompute[subtile_idx], tCompute[subtile_idx + 1])
(
cute.math.exp2(tCompute_log2e[0], fastmath=True),
cute.math.exp2(tCompute_log2e[1], fastmath=True),
),
(tCrDelta[subtile_idx], tCrDelta[subtile_idx + 1]),
)
@@ -3245,11 +3251,11 @@ class SSDKernel:
for reg_idx in range(0, cute.size(tBrB_Compute), 2):
tCompute[reg_idx], tCompute[reg_idx + 1] = cute.arch.mul_packed_f32x2(
(
cute.arch.exp(
(last_column - tBrDeltaA_Compute[reg_idx]).ir_value()
cute.math.exp(
(last_column - tBrDeltaA_Compute[reg_idx]), fastmath=True
),
cute.arch.exp(
(last_column - tBrDeltaA_Compute[reg_idx + 1]).ir_value()
cute.math.exp(
(last_column - tBrDeltaA_Compute[reg_idx + 1]), fastmath=True
),
),
(tBrDelta_Compute[reg_idx], tBrDelta_Compute[reg_idx + 1]),
examples/python/CuTeDSL/hopper/dense_gemm.py
+109 -97
View File
@@ -44,7 +44,7 @@ import cutlass.utils.hopper_helpers as sm90_utils
"""
A high-performance batched dense GEMM (C = A * B) example for the NVIDIA Hopper architecture
using CUTE DSL.
using CuTe DSL.
- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M")
- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K")
- Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M")
@@ -70,7 +70,7 @@ To run this example:
.. code-block:: bash
python examples/hopper/dense_gemm.py \
--mnkl 8192,8192,8192,1 --tile_shape_mnk 128,256,64 \
--mnkl 8192,8192,8192,1 --tile_shape_mn 128,256 \
--cluster_shape_mn 1,1 --a_dtype Float16 --b_dtype Float16 \
--c_dtype Float16 --acc_dtype Float32 \
--a_major k --b_major k --c_major n
@@ -85,7 +85,7 @@ To collect performance with NCU profiler:
.. code-block:: bash
ncu python examples/hopper/dense_gemm.py \
--mnkl 8192,8192,8192,1 --tile_shape_mnk 128,256,64 \
--mnkl 8192,8192,8192,1 --tile_shape_mn 128,256 \
--cluster_shape_mn 1,1 --a_dtype Float16 --b_dtype Float16 \
--c_dtype Float16 --acc_dtype Float32 \
--a_major k --b_major k --c_major n
@@ -95,14 +95,11 @@ Constraints:
* For fp16 types, A and B must have the same data type
* For fp8 types, A and B can have different types (e4m3fn or e5m2) but both must be 8-bit
* Fp8 types only support k-major layout
* Only fp32 accumulation is supported in this example
* CTA tile shape M must be 64/128
* CTA tile shape N must be 64/128/256
* CTA tile shape K must be 64
* Cluster shape M/N must be positive and power of 2, total cluster size <= 4
* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned,
i.e, number of elements is a multiple of 8, 16 for Float16, and Float8, respectively.
* OOB tiles are not allowed when TMA store is disabled
"""
@@ -128,10 +125,10 @@ def parse_arguments() -> argparse.Namespace:
help="mnkl dimensions (comma-separated)",
)
parser.add_argument(
"--tile_shape_mnk",
"--tile_shape_mn",
type=parse_comma_separated_ints,
choices=[(128, 128, 64), (128, 256, 64), (128, 64, 64), (64, 64, 64)],
default=(128, 128, 64),
choices=[(128, 128), (128, 256), (128, 64), (64, 64)],
default=(128, 128),
help="Cta tile shape (comma-separated)",
)
parser.add_argument(
@@ -190,8 +187,8 @@ def parse_arguments() -> argparse.Namespace:
if len(args.mnkl) != 4:
parser.error("--mnkl must contain exactly 4 values")
if len(args.tile_shape_mnk) != 3:
parser.error("--tile_shape_mnk must contain exactly 3 values")
if len(args.tile_shape_mn) != 2:
parser.error("--tile_shape_mn must contain exactly 2 values")
if len(args.cluster_shape_mn) != 2:
parser.error("--cluster_shape_mn must contain exactly 2 values")
@@ -210,10 +207,10 @@ class HopperWgmmaGemmKernel:
:param acc_dtype: Data type for accumulation during computation
:type acc_dtype: type[cutlass.Numeric]
:param tile_shape_mnk: Shape of the CTA tile (M,N,K)
:type tile_shape_mnk: Tuple[int, int, int]
:param cluster_shape_mnk: Cluster dimensions (M,N,K) for parallel processing
:type cluster_shape_mnk: Tuple[int, int, int]
:param tile_shape_mn: Shape of the CTA tile (M,N)
:type tile_shape_mn: Tuple[int, int]
:param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing
:type cluster_shape_mn: Tuple[int, int]
:note: Data type requirements:
- For 16-bit types: A and B must have the same data type
@@ -236,8 +233,8 @@ class HopperWgmmaGemmKernel:
Example:
>>> gemm = HopperWgmmaGemmKernel(
... acc_dtype=cutlass.Float32,
... tile_shape_mnk=(128, 256, 64),
... cluster_shape_mnk=(1, 1, 1)
... tile_shape_mn=(128, 256),
... cluster_shape_mn=(1, 1)
... )
>>> gemm(a_tensor, b_tensor, c_tensor, stream)
"""
@@ -245,8 +242,8 @@ class HopperWgmmaGemmKernel:
def __init__(
self,
acc_dtype: type[cutlass.Numeric],
tile_shape_mnk: tuple[int, int, int],
cluster_shape_mnk: tuple[int, int, int],
tile_shape_mn: tuple[int, int],
cluster_shape_mn: tuple[int, int],
):
"""
Initializes the configuration for a Hopper dense GEMM kernel.
@@ -256,28 +253,30 @@ class HopperWgmmaGemmKernel:
:param acc_dtype: Data type for accumulation during computation
:type acc_dtype: type[cutlass.Numeric]
:param tile_shape_mnk: Shape of the CTA tile (M,N,K)
:type tile_shape_mnk: Tuple[int, int, int]
:param cluster_shape_mnk: Cluster dimensions (M,N,K) for parallel processing
:type cluster_shape_mnk: Tuple[int, int, int]
:param tile_shape_mn: Shape of the CTA tile (M,N)
:type tile_shape_mn: Tuple[int, int]
:param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing
:type cluster_shape_mn: Tuple[int, int]
"""
self.acc_dtype = acc_dtype
self.cluster_shape_mnk = cluster_shape_mnk
self.cluster_shape_mn = cluster_shape_mn
self.mma_inst_shape_mn = None
self.tile_shape_mnk = tuple(tile_shape_mnk)
# K dimension is deferred in _setup_attributes
self.tile_shape_mnk = (*tile_shape_mn, 1)
# For large tile size, using two warp groups is preferred because using only one warp
# group may result in register spill
self.atom_layout_mnk = (
(2, 1, 1)
if tile_shape_mnk[0] > 64 and tile_shape_mnk[1] > 128
if self.tile_shape_mnk[0] > 64 and self.tile_shape_mnk[1] > 128
else (1, 1, 1)
)
self.num_mcast_ctas_a = None
self.num_mcast_ctas_b = None
self.is_a_mcast = False
self.is_b_mcast = False
self.tiled_mma = None
self.occupancy = 1
self.mma_warp_groups = math.prod(self.atom_layout_mnk)
@@ -315,12 +314,27 @@ class HopperWgmmaGemmKernel:
raise ValueError("CTA tile shape M must be 64/128")
if self.tile_shape_mnk[1] not in [64, 128, 256]:
raise ValueError("CTA tile shape N must be 64/128/256")
if self.tile_shape_mnk[2] not in [64]:
raise ValueError("CTA tile shape K must be 64")
self.cta_layout_mnk = cute.make_layout(self.cluster_shape_mnk)
self.num_mcast_ctas_a = self.cluster_shape_mnk[1]
self.num_mcast_ctas_b = self.cluster_shape_mnk[0]
self.tiled_mma = sm90_utils.make_trivial_tiled_mma(
self.a_dtype,
self.b_dtype,
self.a_layout.sm90_mma_major_mode(),
self.b_layout.sm90_mma_major_mode(),
self.acc_dtype,
self.atom_layout_mnk,
tiler_mn=(64, self.tile_shape_mnk[1]),
)
mma_inst_shape_k = cute.size(self.tiled_mma.shape_mnk, mode=[2])
mma_inst_tile_k = 4
self.tile_shape_mnk = (
self.tile_shape_mnk[0],
self.tile_shape_mnk[1],
mma_inst_shape_k * mma_inst_tile_k,
)
self.cta_layout_mnk = cute.make_layout((*self.cluster_shape_mn, 1))
self.num_mcast_ctas_a = self.cluster_shape_mn[1]
self.num_mcast_ctas_b = self.cluster_shape_mn[0]
self.is_a_mcast = self.num_mcast_ctas_a > 1
self.is_b_mcast = self.num_mcast_ctas_b > 1
@@ -401,28 +415,18 @@ class HopperWgmmaGemmKernel:
self._setup_attributes()
tiled_mma = sm90_utils.make_trivial_tiled_mma(
self.a_dtype,
self.b_dtype,
self.a_layout.sm90_mma_major_mode(),
self.b_layout.sm90_mma_major_mode(),
self.acc_dtype,
self.atom_layout_mnk,
tiler_mn=(64, self.tile_shape_mnk[1]),
)
tma_atom_a, tma_tensor_a = self._make_tma_atoms_and_tensors(
a,
self.a_smem_layout_staged,
(self.tile_shape_mnk[0], self.tile_shape_mnk[2]),
self.cluster_shape_mnk[1],
self.cluster_shape_mn[1],
)
tma_atom_b, tma_tensor_b = self._make_tma_atoms_and_tensors(
b,
self.b_smem_layout_staged,
(self.tile_shape_mnk[1], self.tile_shape_mnk[2]),
self.cluster_shape_mnk[0],
self.cluster_shape_mn[0],
)
tma_atom_c, tma_tensor_c = self._make_tma_store_atoms_and_tensors(
@@ -431,20 +435,20 @@ class HopperWgmmaGemmKernel:
self.epi_tile,
)
grid = self._compute_grid(c, self.tile_shape_mnk, self.cluster_shape_mnk)
grid = self._compute_grid(c, self.tile_shape_mnk, self.cluster_shape_mn)
@cute.struct
class SharedStorage:
mainloop_pipeline_array_ptr: cute.struct.MemRange[
cutlass.Int64, self.ab_stage * 2
]
sa: cute.struct.Align[
sA: cute.struct.Align[
cute.struct.MemRange[
self.a_dtype, cute.cosize(self.a_smem_layout_staged)
],
self.buffer_align_bytes,
]
sb: cute.struct.Align[
sB: cute.struct.Align[
cute.struct.MemRange[
self.b_dtype, cute.cosize(self.b_smem_layout_staged)
],
@@ -461,7 +465,7 @@ class HopperWgmmaGemmKernel:
tma_tensor_b,
tma_atom_c,
tma_tensor_c,
tiled_mma,
self.tiled_mma,
self.cta_layout_mnk,
self.a_smem_layout_staged,
self.b_smem_layout_staged,
@@ -469,8 +473,7 @@ class HopperWgmmaGemmKernel:
).launch(
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=self.cluster_shape_mnk,
smem=self.shared_storage.size_in_bytes(),
cluster=(*self.cluster_shape_mn, 1),
stream=stream,
)
return
@@ -562,8 +565,8 @@ class HopperWgmmaGemmKernel:
# Get the pid from cluster id
bidx_in_cluster = cute.arch.block_in_cluster_idx()
pid_m = cid_m * self.cluster_shape_mnk[0] + bidx_in_cluster[0]
pid_n = cid_n * self.cluster_shape_mnk[1] + bidx_in_cluster[1]
pid_m = cid_m * self.cluster_shape_mn[0] + bidx_in_cluster[0]
pid_n = cid_n * self.cluster_shape_mn[1] + bidx_in_cluster[1]
tile_coord_mnkl = (pid_m, pid_n, None, bidz)
cta_rank_in_cluster = cute.arch.make_warp_uniform(
@@ -621,22 +624,22 @@ class HopperWgmmaGemmKernel:
)
# Cluster arrive after barrier init
if cute.size(self.cluster_shape_mnk) > 1:
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_arrive_relaxed()
# ///////////////////////////////////////////////////////////////////////////////
# Generate smem tensor A/B
# ///////////////////////////////////////////////////////////////////////////////
sa = storage.sa.get_tensor(
sA = storage.sA.get_tensor(
a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner
)
sb = storage.sb.get_tensor(
sB = storage.sB.get_tensor(
b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner
)
sc_ptr = cute.recast_ptr(
sa.iterator, epi_smem_layout_staged.inner, dtype=self.c_dtype
sC_ptr = cute.recast_ptr(
sA.iterator, epi_smem_layout_staged.inner, dtype=self.c_dtype
)
sc = cute.make_tensor(sc_ptr, epi_smem_layout_staged.outer)
sC = cute.make_tensor(sC_ptr, epi_smem_layout_staged.outer)
# ///////////////////////////////////////////////////////////////////////////////
# Local_tile partition global tensors
@@ -673,34 +676,34 @@ class HopperWgmmaGemmKernel:
# TMA load A partition_S/D
a_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (0, None, 0)).shape)
a_cta_crd = cluster_coord_mnk[1]
sa_for_tma_partition = cute.group_modes(sa, 0, 2)
sA_for_tma_partition = cute.group_modes(sA, 0, 2)
gA_for_tma_partition = cute.group_modes(gA_mkl, 0, 2)
tAsA, tAgA_mkl = cute.nvgpu.cpasync.tma_partition(
tma_atom_a,
a_cta_crd,
a_cta_layout,
sa_for_tma_partition,
sA_for_tma_partition,
gA_for_tma_partition,
)
# TMA load B partition_S/D
b_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (None, 0, 0)).shape)
b_cta_crd = cluster_coord_mnk[0]
sb_for_tma_partition = cute.group_modes(sb, 0, 2)
sB_for_tma_partition = cute.group_modes(sB, 0, 2)
gB_for_tma_partition = cute.group_modes(gB_nkl, 0, 2)
tBsB, tBgB_nkl = cute.nvgpu.cpasync.tma_partition(
tma_atom_b,
b_cta_crd,
b_cta_layout,
sb_for_tma_partition,
sB_for_tma_partition,
gB_for_tma_partition,
)
# //////////////////////////////////////////////////////////////////////////////
# Make frangments
# Make fragments
# //////////////////////////////////////////////////////////////////////////////
tCsA = thr_mma.partition_A(sa)
tCsB = thr_mma.partition_B(sb)
tCsA = thr_mma.partition_A(sA)
tCsB = thr_mma.partition_B(sB)
tCrA = tiled_mma.make_fragment_A(tCsA)
tCrB = tiled_mma.make_fragment_B(tCsB)
@@ -711,7 +714,7 @@ class HopperWgmmaGemmKernel:
# Cluster wait
# ///////////////////////////////////////////////////////////////////////////////
# cluster wait for barrier init
if cute.size(self.cluster_shape_mnk) > 1:
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
else:
cute.arch.sync_threads()
@@ -788,7 +791,7 @@ class HopperWgmmaGemmKernel:
tiled_mma.set(cute.nvgpu.warpgroup.Field.ACCUMULATE, False)
num_k_blocks = cute.size(tCrA, mode=[2])
for k_tile in range(k_pipe_mmas):
for k_tile in cutlass.range_constexpr(k_pipe_mmas):
# Wait for A/B buffer to be ready
mainloop_pipeline.consumer_wait(
mainloop_consumer_read_state, peek_ab_full_status
@@ -917,7 +920,7 @@ class HopperWgmmaGemmKernel:
# /////////////////////////////////////////////////////////////////////////////
cute.nvgpu.warpgroup.wait_group(0)
if cute.size(self.cluster_shape_mnk) > 1:
if cute.size(self.cluster_shape_mn) > 1:
# Wait for all threads in the cluster to finish, avoid early release of smem
cute.arch.cluster_arrive()
cute.arch.cluster_wait()
@@ -950,33 +953,45 @@ class HopperWgmmaGemmKernel:
# (R2S, R2S_M, R2S_N, PIPE_D)
thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
tRS_sD = thr_copy_r2s.partition_D(sc)
tRS_sD = thr_copy_r2s.partition_D(sC)
# (R2S, R2S_M, R2S_N)
tRS_rAcc = tiled_copy_r2s.retile(accumulators)
# Allocate D registers.
rD_shape = cute.shape(thr_copy_r2s.partition_S(sc))
rD_shape = cute.shape(thr_copy_r2s.partition_S(sC))
tRS_rD_layout = cute.make_layout(rD_shape[:3])
tRS_rD = cute.make_fragment_like(tRS_rD_layout, self.acc_dtype)
size_tRS_rD = cute.size(tRS_rD)
sepi_for_tma_partition = cute.group_modes(sc, 0, 2)
tcgc_for_tma_partition = cute.zipped_divide(gC_mnl, self.epi_tile)
sepi_for_tma_partition = cute.group_modes(sC, 0, 2)
tCgC_for_tma_partition = cute.zipped_divide(gC_mnl, self.epi_tile)
bSG_sD, bSG_gD = cute.nvgpu.cpasync.tma_partition(
tma_atom_c,
0,
cute.make_layout(1),
sepi_for_tma_partition,
tcgc_for_tma_partition,
tCgC_for_tma_partition,
)
epi_tile_num = cute.size(tcgc_for_tma_partition, mode=[1])
epi_tile_shape = tcgc_for_tma_partition.shape[1]
epi_tile_num = cute.size(tCgC_for_tma_partition, mode=[1])
epi_tile_shape = tCgC_for_tma_partition.shape[1]
epi_tile_layout = cute.make_layout(
epi_tile_shape, stride=(epi_tile_shape[1], 1)
)
for epi_idx in cutlass.range(epi_tile_num, unroll=epi_tile_num):
# Initialize tma store c_pipeline
c_producer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, self.threads_per_cta, self.threads_per_cta
)
c_pipeline = pipeline.PipelineTmaStore.create(
num_stages=self.epi_stage,
producer_group=c_producer_group,
)
for epi_idx in cutlass.range_constexpr(epi_tile_num):
# Copy from accumulators to D registers
for epi_v in range(size_tRS_rD):
for epi_v in cutlass.range_constexpr(size_tRS_rD):
tRS_rD[epi_v] = tRS_rAcc[epi_idx * size_tRS_rD + epi_v]
# Type conversion
@@ -997,10 +1012,6 @@ class HopperWgmmaGemmKernel:
# barrier for sync
cute.arch.barrier()
# Get the global memory coordinate for the current epi tile.
epi_tile_layout = cute.make_layout(
epi_tile_shape, stride=(epi_tile_shape[1], 1)
)
gmem_coord = epi_tile_layout.get_hier_coord(epi_idx)
# Copy from shared memory to global memory
if warp_idx == 0:
@@ -1009,11 +1020,14 @@ class HopperWgmmaGemmKernel:
bSG_sD[(None, epi_buffer)],
bSG_gD[(None, gmem_coord)],
)
cute.arch.cp_async_bulk_commit_group()
cute.arch.cp_async_bulk_wait_group(self.epi_stage - 1, read=True)
c_pipeline.producer_commit()
c_pipeline.producer_acquire()
cute.arch.barrier()
if warp_idx == 0:
c_pipeline.producer_tail()
return
@staticmethod
@@ -1055,9 +1069,7 @@ class HopperWgmmaGemmKernel:
mbar_helpers_bytes = 1024
ab_stage = (
(smem_capacity - occupancy * 1024) // occupancy
- mbar_helpers_bytes
- epi_bytes
smem_capacity // occupancy - mbar_helpers_bytes - epi_bytes
) // ab_bytes_per_stage
return ab_stage, epi_stage
@@ -1195,7 +1207,7 @@ class HopperWgmmaGemmKernel:
def _compute_grid(
c: cute.Tensor,
tile_shape_mnk: tuple[int, int, int],
cluster_shape_mnk: tuple[int, int, int],
cluster_shape_mn: tuple[int, int],
) -> tuple[int, int, int]:
"""Compute grid shape for the output tensor C.
@@ -1203,8 +1215,8 @@ class HopperWgmmaGemmKernel:
:type c: cute.Tensor
:param tile_shape_mnk: The shape (M, N, K) of the CTA tile.
:type tile_shape_mnk: tuple[int, int, int]
:param cluster_shape_mnk: Shape of each cluster in M, N, K dimensions.
:type cluster_shape_mnk: tuple[int, int, int]
:param cluster_shape_mn: Shape of each cluster in M, N dimensions.
:type cluster_shape_mn: tuple[int, int]
:return: Grid shape for kernel launch.
:rtype: tuple[int, int, int]
@@ -1212,8 +1224,9 @@ class HopperWgmmaGemmKernel:
c_shape = (tile_shape_mnk[0], tile_shape_mnk[1])
gc = cute.zipped_divide(c, tiler=c_shape)
clusters = cute.ceil_div(cute.get(gc.layout, mode=[1]).shape, cluster_shape_mnk)
grid = tuple(x * y for x, y in zip(clusters, cluster_shape_mnk))
cluster_shape_mnl = (*cluster_shape_mn, 1)
clusters = cute.ceil_div(cute.get(gc.layout, mode=[1]).shape, cluster_shape_mnl)
grid = tuple(x * y for x, y in zip(clusters, cluster_shape_mnl))
return grid
@staticmethod
@@ -1363,7 +1376,7 @@ def run(
a_major: str,
b_major: str,
c_major: str,
tile_shape_mnk: Tuple[int, int, int],
tile_shape_mn: Tuple[int, int],
cluster_shape_mn: Tuple[int, int],
tolerance: float,
warmup_iterations: int,
@@ -1387,8 +1400,8 @@ def run(
: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 tile_shape_mnk: CTA tile shape (M, N, K)
:type tile_shape_mnk: Tuple[int, int, int]
:param tile_shape_mn: CTA tile shape (M, N)
:type tile_shape_mn: Tuple[int, int]
:param cluster_shape_mn: Cluster shape (M, N)
:type cluster_shape_mn: Tuple[int, int]
:param tolerance: Tolerance value for reference validation comparison
@@ -1411,7 +1424,7 @@ def run(
f"A dtype: {a_dtype}, B dtype: {b_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}"
)
print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}")
print(f"Tile Shape: {tile_shape_mnk}, Cluster Shape: {cluster_shape_mn}")
print(f"Tile Shape: {tile_shape_mn}, Cluster Shape: {cluster_shape_mn}")
print(f"Tolerance: {tolerance}")
print(f"Warmup iterations: {warmup_iterations}")
print(f"Iterations: {iterations}")
@@ -1420,7 +1433,6 @@ def run(
# Unpack parameters
m, n, k, l = mnkl
cluster_shape_mnk = (*cluster_shape_mn, 1)
# Skip unsupported types
if not HopperWgmmaGemmKernel.is_valid_dtypes(
@@ -1488,7 +1500,7 @@ def run(
b, mB, b_torch = create_and_permute_tensor(l, n, k, b_major == "n", b_dtype)
c, mC, c_torch = create_and_permute_tensor(l, m, n, c_major == "m", c_dtype)
gemm = HopperWgmmaGemmKernel(acc_dtype, tile_shape_mnk, cluster_shape_mnk)
gemm = HopperWgmmaGemmKernel(acc_dtype, tile_shape_mn, cluster_shape_mn)
torch_stream = torch.cuda.Stream()
stream = cuda.CUstream(torch_stream.cuda_stream)
@@ -1572,7 +1584,7 @@ if __name__ == "__main__":
args.a_major,
args.b_major,
args.c_major,
args.tile_shape_mnk,
args.tile_shape_mn,
args.cluster_shape_mn,
args.tolerance,
args.warmup_iterations,
examples/python/CuTeDSL/notebooks/tensorssa.ipynb
+151 -207
View File
@@ -2,7 +2,7 @@
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -41,22 +41,9 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"a_vec: tensor_value<vector<12xf32> o (3, 4)>\n",
"b_vec: tensor_value<vector<12xf32> o (3, 4)>\n",
"tensor(raw_ptr(0x0000000006cff170: f32, generic, align<4>) o (3,4):(4,1), data=\n",
" [[ 2.000000, 2.000000, 2.000000, 2.000000, ],\n",
" [ 2.000000, 2.000000, 2.000000, 2.000000, ],\n",
" [ 2.000000, 2.000000, 2.000000, 2.000000, ]])\n"
]
}
],
"outputs": [],
"source": [
"@cute.jit\n",
"def load_and_store(res: cute.Tensor, a: cute.Tensor, b: cute.Tensor):\n",
@@ -91,22 +78,9 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor_value<vector<24xf32> o (4, 2, 3)> -> tensor_value<vector<12xf32> o (4, 3)>\n",
"tensor(raw_ptr(0x00000000071acaf0: f32, generic, align<4>) o (4,3):(3,1), data=\n",
" [[ 3.000000, 4.000000, 5.000000, ],\n",
" [ 9.000000, 10.000000, 11.000000, ],\n",
" [ 15.000000, 16.000000, 17.000000, ],\n",
" [ 21.000000, 22.000000, 23.000000, ]])\n"
]
}
],
"outputs": [],
"source": [
"@cute.jit\n",
"def apply_slice(src: cute.Tensor, dst: cute.Tensor, indices: cutlass.Constexpr):\n",
@@ -155,19 +129,9 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor_value<vector<24xf32> o (4, 2, 3)> -> ?\n",
"tensor(raw_ptr(0x00000000013cbbe0: f32, generic, align<4>) o (1):(1), data=\n",
" [ 10.000000, ])\n"
]
}
],
"outputs": [],
"source": [
"def slice_2():\n",
" src_shape = (4, 2, 3)\n",
@@ -195,40 +159,9 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor(raw_ptr(0x00000000074f0e70: f32, generic, align<4>) o (3):(1), data=\n",
" [ 3.000000, ],\n",
" [ 3.000000, ],\n",
" [ 3.000000, ])\n",
"tensor(raw_ptr(0x00000000074f0e70: f32, generic, align<4>) o (3):(1), data=\n",
" [-1.000000, ],\n",
" [-1.000000, ],\n",
" [-1.000000, ])\n",
"tensor(raw_ptr(0x00000000074f0e70: f32, generic, align<4>) o (3):(1), data=\n",
" [ 2.000000, ],\n",
" [ 2.000000, ],\n",
" [ 2.000000, ])\n",
"tensor(raw_ptr(0x00000000074f0e70: f32, generic, align<4>) o (3):(1), data=\n",
" [ 0.500000, ],\n",
" [ 0.500000, ],\n",
" [ 0.500000, ])\n",
"tensor(raw_ptr(0x00000000074f0e70: f32, generic, align<4>) o (3):(1), data=\n",
" [ 0.000000, ],\n",
" [ 0.000000, ],\n",
" [ 0.000000, ])\n",
"tensor(raw_ptr(0x00000000074f0e70: f32, generic, align<4>) o (3):(1), data=\n",
" [ 1.000000, ],\n",
" [ 1.000000, ],\n",
" [ 1.000000, ])\n"
]
}
],
"outputs": [],
"source": [
"@cute.jit\n",
"def binary_op_1(res: cute.Tensor, a: cute.Tensor, b: cute.Tensor):\n",
@@ -236,28 +169,22 @@
" b_vec = b.load()\n",
"\n",
" add_res = a_vec + b_vec\n",
" res.store(add_res)\n",
" cute.print_tensor(res) # prints [3.000000, 3.000000, 3.000000]\n",
" cute.print_tensor(add_res) # prints [3.000000, 3.000000, 3.000000]\n",
"\n",
" sub_res = a_vec - b_vec\n",
" res.store(sub_res)\n",
" cute.print_tensor(res) # prints [-1.000000, -1.000000, -1.000000]\n",
" cute.print_tensor(sub_res) # prints [-1.000000, -1.000000, -1.000000]\n",
"\n",
" mul_res = a_vec * b_vec\n",
" res.store(mul_res)\n",
" cute.print_tensor(res) # prints [2.000000, 2.000000, 2.000000]\n",
" cute.print_tensor(mul_res) # prints [2.000000, 2.000000, 2.000000]\n",
"\n",
" div_res = a_vec / b_vec\n",
" res.store(div_res)\n",
" cute.print_tensor(res) # prints [0.500000, 0.500000, 0.500000]\n",
" cute.print_tensor(div_res) # prints [0.500000, 0.500000, 0.500000]\n",
"\n",
" floor_div_res = a_vec // b_vec\n",
" res.store(floor_div_res)\n",
" cute.print_tensor(res) # prints [0.000000, 0.000000, 0.000000]\n",
" cute.print_tensor(res) # prints [0.000000, 0.000000, 0.000000]\n",
"\n",
" mod_res = a_vec % b_vec\n",
" res.store(mod_res)\n",
" cute.print_tensor(res) # prints [1.000000, 1.000000, 1.000000]\n",
" cute.print_tensor(mod_res) # prints [1.000000, 1.000000, 1.000000]\n",
"\n",
"\n",
"a = np.empty((3,), dtype=np.float32)\n",
@@ -270,68 +197,31 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor(raw_ptr(0x0000000007828ed0: f32, generic, align<4>) o (3):(1), data=\n",
" [ 3.000000, ],\n",
" [ 3.000000, ],\n",
" [ 3.000000, ])\n",
"tensor(raw_ptr(0x0000000007828ed0: f32, generic, align<4>) o (3):(1), data=\n",
" [-1.000000, ],\n",
" [-1.000000, ],\n",
" [-1.000000, ])\n",
"tensor(raw_ptr(0x0000000007828ed0: f32, generic, align<4>) o (3):(1), data=\n",
" [ 2.000000, ],\n",
" [ 2.000000, ],\n",
" [ 2.000000, ])\n",
"tensor(raw_ptr(0x0000000007828ed0: f32, generic, align<4>) o (3):(1), data=\n",
" [ 0.500000, ],\n",
" [ 0.500000, ],\n",
" [ 0.500000, ])\n",
"tensor(raw_ptr(0x0000000007828ed0: f32, generic, align<4>) o (3):(1), data=\n",
" [ 0.000000, ],\n",
" [ 0.000000, ],\n",
" [ 0.000000, ])\n",
"tensor(raw_ptr(0x0000000007828ed0: f32, generic, align<4>) o (3):(1), data=\n",
" [ 1.000000, ],\n",
" [ 1.000000, ],\n",
" [ 1.000000, ])\n"
]
}
],
"outputs": [],
"source": [
"@cute.jit\n",
"def binary_op_2(res: cute.Tensor, a: cute.Tensor, c: cutlass.Constexpr):\n",
" a_vec = a.load()\n",
"\n",
" add_res = a_vec + c\n",
" res.store(add_res)\n",
" cute.print_tensor(res) # prints [3.000000, 3.000000, 3.000000]\n",
" cute.print_tensor(add_res) # prints [3.000000, 3.000000, 3.000000]\n",
"\n",
" sub_res = a_vec - c\n",
" res.store(sub_res)\n",
" cute.print_tensor(res) # prints [-1.000000, -1.000000, -1.000000]\n",
" cute.print_tensor(sub_res) # prints [-1.000000, -1.000000, -1.000000]\n",
"\n",
" mul_res = a_vec * c\n",
" res.store(mul_res)\n",
" cute.print_tensor(res) # prints [2.000000, 2.000000, 2.000000]\n",
" cute.print_tensor(mul_res) # prints [2.000000, 2.000000, 2.000000]\n",
"\n",
" div_res = a_vec / c\n",
" res.store(div_res)\n",
" cute.print_tensor(res) # prints [0.500000, 0.500000, 0.500000]\n",
" cute.print_tensor(div_res) # prints [0.500000, 0.500000, 0.500000]\n",
"\n",
" floor_div_res = a_vec // c\n",
" res.store(floor_div_res)\n",
" cute.print_tensor(res) # prints [0.000000, 0.000000, 0.000000]\n",
" cute.print_tensor(floor_div_res) # prints [0.000000, 0.000000, 0.000000]\n",
"\n",
" mod_res = a_vec % c\n",
" res.store(mod_res)\n",
" cute.print_tensor(res) # prints [1.000000, 1.000000, 1.000000]\n",
" cute.print_tensor(mod_res) # prints [1.000000, 1.000000, 1.000000]\n",
"\n",
"a = np.empty((3,), dtype=np.float32)\n",
"a.fill(1.0)\n",
@@ -342,17 +232,9 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[False True False]\n"
]
}
],
"outputs": [],
"source": [
"@cute.jit\n",
"def binary_op_3(res: cute.Tensor, a: cute.Tensor, b: cute.Tensor):\n",
@@ -378,17 +260,9 @@
},
{
"cell_type": "code",
"execution_count": 8,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[3 0 7]\n"
]
}
],
"outputs": [],
"source": [
"@cute.jit\n",
"def binary_op_4(res: cute.Tensor, a: cute.Tensor, b: cute.Tensor):\n",
@@ -420,44 +294,23 @@
},
{
"cell_type": "code",
"execution_count": 9,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor(raw_ptr(0x0000000007fbd180: f32, generic, align<4>) o (3):(1), data=\n",
" [ 2.000000, ],\n",
" [ 2.000000, ],\n",
" [ 2.000000, ])\n",
"tensor(raw_ptr(0x0000000007fbd180: f32, generic, align<4>) o (3):(1), data=\n",
" [-0.756802, ],\n",
" [-0.756802, ],\n",
" [-0.756802, ])\n",
"tensor(raw_ptr(0x0000000007fbd180: f32, generic, align<4>) o (3):(1), data=\n",
" [ 16.000000, ],\n",
" [ 16.000000, ],\n",
" [ 16.000000, ])\n"
]
}
],
"outputs": [],
"source": [
"@cute.jit\n",
"def unary_op_1(res: cute.Tensor, a: cute.Tensor):\n",
" a_vec = a.load()\n",
"\n",
" sqrt_res = cute.math.sqrt(a_vec)\n",
" res.store(sqrt_res)\n",
" cute.print_tensor(res) # prints [2.000000, 2.000000, 2.000000]\n",
" cute.print_tensor(sqrt_res) # prints [2.000000, 2.000000, 2.000000]\n",
"\n",
" sin_res = cute.math.sin(a_vec)\n",
" res.store(sin_res)\n",
" cute.print_tensor(res) # prints [-0.756802, -0.756802, -0.756802]\n",
" cute.print_tensor(sin_res) # prints [-0.756802, -0.756802, -0.756802]\n",
"\n",
" exp2_res = cute.math.exp2(a_vec)\n",
" res.store(exp2_res)\n",
" cute.print_tensor(res) # prints [16.000000, 16.000000, 16.000000]\n",
" cute.print_tensor(exp2_res) # prints [16.000000, 16.000000, 16.000000]\n",
"\n",
"a = np.array([4.0, 4.0, 4.0], dtype=np.float32)\n",
"res = np.empty((3,), dtype=np.float32)\n",
@@ -470,29 +323,18 @@
"source": [
"#### Reduction Operation\n",
"\n",
"The `TensorSSA`'s `reduce` method applies a specified reduction operation (`ReductionOp.ADD`, `ReductionOp.MUL`, `ReductionOp.MAX`, `ReductionOp.MIN`) starting with an initial value, and performs this reduction along the dimensions specified by the `reduction_profile.`. The result is typically a new `TensorSSA` with reduced dimensions or a scalar value if reduces across all axes."
"The `TensorSSA`'s `reduce` method applies a specified reduction operation (`ReductionOp.ADD`, \n",
"`ReductionOp.MUL`, `ReductionOp.MAX`, `ReductionOp.MIN`) starting with an initial value, and \n",
"performs this reduction along the dimensions specified by the `reduction_profile`. The result \n",
"is typically a new `TensorSSA` with reduced dimensions or a scalar value if it reduces across \n",
"all axes."
]
},
{
"cell_type": "code",
"execution_count": 10,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"21.000000\n",
"tensor(raw_ptr(0x00007ffd1ea2bca0: f32, rmem, align<32>) o (2):(1), data=\n",
" [ 6.000000, ],\n",
" [ 15.000000, ])\n",
"tensor(raw_ptr(0x00007ffd1ea2bcc0: f32, rmem, align<32>) o (3):(1), data=\n",
" [ 6.000000, ],\n",
" [ 8.000000, ],\n",
" [ 10.000000, ])\n"
]
}
],
"outputs": [],
"source": [
"@cute.jit\n",
"def reduction_op(a: cute.Tensor):\n",
@@ -507,36 +349,138 @@
" 0.0,\n",
" reduction_profile=0\n",
" )\n",
" cute.printf(red_res) # prints 21.000000\n",
" cute.printf(red_res) # prints 21.000000\n",
"\n",
" red_res = a_vec.reduce(\n",
" cute.ReductionOp.ADD,\n",
" 0.0,\n",
" reduction_profile=(None, 1)\n",
" )\n",
" # We can't print the TensorSSA directly at this point, so we store it to a new Tensor and print it.\n",
" res = cute.make_fragment(red_res.shape, cutlass.Float32)\n",
" res.store(red_res)\n",
" cute.print_tensor(res) # prints [6.000000, 15.000000]\n",
" cute.print_tensor(red_res) # prints [6.000000, 15.000000]\n",
"\n",
" red_res = a_vec.reduce(\n",
" cute.ReductionOp.ADD,\n",
" 1.0,\n",
" reduction_profile=(1, None)\n",
" )\n",
" res = cute.make_fragment(red_res.shape, cutlass.Float32)\n",
" res.store(red_res)\n",
" cute.print_tensor(res) # prints [6.000000, 8.000000, 10.000000]\n",
" cute.print_tensor(red_res) # prints [6.000000, 8.000000, 10.000000]\n",
"\n",
"\n",
"a = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)\n",
"reduction_op(from_dlpack(a))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Broadcast\n",
"\n",
"`TensorSSA` supports broadcasting operations following NumPy's broadcasting rules. Broadcasting \n",
"allows you to perform operations on arrays of different shapes when certain conditions are met. \n",
"The key rules are:\n",
"\n",
"1. Source shape is padded with 1's to match the rank of target shape\n",
"2. The size in each mode of source shape must either be 1 or equal to target shape\n",
"3. After broadcasting, all modes should match target shape\n",
"\n",
"Let's look at some examples of broadcasting in action:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import cutlass\n",
"import cutlass.cute as cute\n",
"\n",
"\n",
"@cute.jit\n",
"def broadcast_examples():\n",
" a = cute.make_fragment((1,3), dtype=cutlass.Float32)\n",
" a[0] = 0.0\n",
" a[1] = 1.0\n",
" a[2] = 2.0\n",
" a_val = a.load()\n",
" cute.print_tensor(a_val.broadcast_to((4, 3)))\n",
" # tensor(raw_ptr(0x00007ffe26625740: f32, rmem, align<32>) o (4,3):(1,4), data=\n",
" # [[ 0.000000, 1.000000, 2.000000, ],\n",
" # [ 0.000000, 1.000000, 2.000000, ],\n",
" # [ 0.000000, 1.000000, 2.000000, ],\n",
" # [ 0.000000, 1.000000, 2.000000, ]])\n",
"\n",
" c = cute.make_fragment((4,1), dtype=cutlass.Float32)\n",
" c[0] = 0.0\n",
" c[1] = 1.0\n",
" c[2] = 2.0\n",
" c[3] = 3.0\n",
" cute.print_tensor(a.load() + c.load())\n",
" # tensor(raw_ptr(0x00007ffe26625780: f32, rmem, align<32>) o (4,3):(1,4), data=\n",
" # [[ 0.000000, 1.000000, 2.000000, ],\n",
" # [ 1.000000, 2.000000, 3.000000, ],\n",
" # [ 2.000000, 3.000000, 4.000000, ],\n",
" # [ 3.000000, 4.000000, 5.000000, ]])\n",
"\n",
"\n",
"broadcast_examples()"
]
},
{
"cell_type": "markdown",
"metadata": {
"vscode": {
"languageId": "raw"
}
},
"source": [
"The examples above demonstrate two key broadcasting scenarios:\n",
"\n",
"1. **Row Vector Broadcasting**: In the first example, we create a row vector `a` with shape \n",
" (1, 3) containing values [0.0, 1.0, 2.0]. When we broadcast it to shape (4, 3), the values \n",
" are repeated across the first dimension, resulting in:\n",
" ```\n",
" [[0.0, 1.0, 2.0],\n",
" [0.0, 1.0, 2.0],\n",
" [0.0, 1.0, 2.0],\n",
" [0.0, 1.0, 2.0]]\n",
" ```\n",
" This demonstrates how a row vector can be broadcast to create multiple identical rows.\n",
"\n",
"2. **Column Vector and Row Vector Addition**: In the second example, we have:\n",
" - A row vector `a` with shape (1, 3) containing [0.0, 1.0, 2.0]\n",
" - A column vector `c` with shape (4, 1) containing [0.0, 1.0, 2.0, 3.0]\n",
" \n",
" When we add these together, both vectors are broadcast to shape (4, 3):\n",
" - The row vector is broadcast vertically (4 times)\n",
" - The column vector is broadcast horizontally (3 times)\n",
" \n",
" The result is:\n",
" ```\n",
" [[0.0 + 0.0, 1.0 + 0.0, 2.0 + 0.0],\n",
" [0.0 + 1.0, 1.0 + 1.0, 2.0 + 1.0],\n",
" [0.0 + 2.0, 1.0 + 2.0, 2.0 + 2.0],\n",
" [0.0 + 3.0, 1.0 + 3.0, 2.0 + 3.0]]\n",
" ```\n",
" =\n",
" ```\n",
" [[0.0, 1.0, 2.0],\n",
" [1.0, 2.0, 3.0],\n",
" [2.0, 3.0, 4.0],\n",
" [3.0, 4.0, 5.0]]\n",
" ```\n",
"\n",
"This demonstrates how `TensorSSA` can automatically handle broadcasting of both row and column \n",
"vectors in arithmetic operations, following the broadcasting rules where each dimension must \n",
"either be 1 or match the target size. The broadcasting is handled implicitly during operations, \n",
"making it easy to work with tensors of different shapes.\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"display_name": ".venv3_12",
"language": "python",
"name": "python3"
},
@@ -550,7 +494,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.5"
"version": "3.12.10"
}
},
"nbformat": 4,
include/cute/atom/copy_atom.hpp
+2 -1
View File
@@ -539,7 +539,8 @@ make_cotiled_copy(Copy_Atom<Args...> const& copy_atom,
auto layout_tv_data = composition(inv_data_layout, atom_tv_layout);
// Check validity
CUTE_STATIC_ASSERT_V(coalesce(composition(data_layout, layout<1>(layout_tv_data))) == coalesce(layout<1>(atom_tv_layout)),
// Append 1:0 to data_layout so that OOB coordinates get the stride-0
CUTE_STATIC_ASSERT_V(coalesce(composition(make_layout(data_layout, Layout<_1,_0>{}), layout<1>(layout_tv_data))) == coalesce(layout<1>(atom_tv_layout)),
"The memory pointed to by AtomTVLayout does not exist in the DataLayout.");
//
// Tiler -- Find the active elements in the DATA tensor and generate a tiler to extract them
include/cutlass/detail/collective/mixed_input_utils.hpp
+3 -21
View File
@@ -705,7 +705,7 @@ public:
auto smem_tiled_copy_S = cute::get<0>(partitioned_transform_extra_info);
auto&& scales = cute::get<1>(partitioned_transform_extra_info);
using ScaleType = decltype(scales);
auto tSrS = make_tensor(static_cast<ScaleType&&>(scales).data(), scales.layout());
auto tSrS = make_tensor(scales.data(), scales.layout());
auto tSsS = cute::get<2>(partitioned_transform_extra_info);
copy(smem_tiled_copy_S, tSsS(_,_,_,_,load2transform_consumer_index), tSrS);
@@ -714,7 +714,7 @@ public:
} else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
auto&& zeros = cute::get<3>(partitioned_transform_extra_info);
using ZeroType = decltype(zeros);
auto tZrZ = make_tensor(static_cast<ZeroType&&>(zeros).data(), zeros.layout());
auto tZrZ = make_tensor(zeros.data(), zeros.layout());
auto tZsZ = cute::get<4>(partitioned_transform_extra_info);
copy(smem_tiled_copy_S, tZsZ(_,_,_,_,load2transform_consumer_index), tZrZ);
@@ -1002,7 +1002,6 @@ public:
auto src_arr = recast<SrcArray>(src);
auto dst_arr = recast<DstArray>(dst);
Tensor src_vm = cute::group_modes<1,-1>(cute::zipped_divide(src, pack));
Tensor dst_vm = cute::group_modes<1,-1>(cute::zipped_divide(dst, pack));
cute::transform(src_arr, dst_arr, Converter::convert);
@@ -1019,7 +1018,6 @@ public:
auto scale_arr = recast<ScaleArray>(filter_zeros(scales));
if constexpr (is_same_v<DstType, cutlass::bfloat16_t>){
Tensor dst_vm = cute::group_modes<1,-1>(cute::zipped_divide(dst, pack));
Tensor scales_vm = cute::group_modes<1,-1>(cute::zipped_divide(scales, pack));
for (int i = 0; i < size<1>(dst_vm); ++i){
@@ -1194,13 +1192,7 @@ public:
Tensor tCsS = cta_mma.partition_A(sS);
Tensor tSsS = smem_thr_copy_S.partition_S(tCsS);
Tensor tSrS = make_tensor<ElementScale>(tSsS(_,_,_,_,0).shape());
#if 0
if(cute::thread(128, 0)){
print("sS: ");print(sS);print("\n");
print("tSsS: ");print(tSsS);print("\n");
print("tSrS: ");print(tSrS);print("\n");
}
#endif
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
return cute::make_tuple(smem_tiled_copy_S, tSrS, tSsS);
}
@@ -1209,16 +1201,6 @@ public:
Tensor tCsZ = cta_mma.partition_A(sZ);
Tensor tZsZ = smem_thr_copy_S.partition_S(tCsZ);
Tensor tZrZ = make_tensor<ElementZero>(tZsZ(_,_,_,_,0).shape());
#if 0
if(cute::thread(128, 0)){
print("sS: ");print(sS);print("\n");
print("tSsS: ");print(tSsS);print("\n");
print("tSrS: ");print(tSrS);print("\n");
print("sZ: ");print(sZ);print("\n");
print("tZsZ: ");print(tZsZ);print("\n");
print("tZrZ: ");print(tZrZ);print("\n");
}
#endif
return cute::make_tuple(smem_tiled_copy_S, tSrS, tSsS, tZrZ, tZsZ);
}
else {
include/cutlass/epilogue/collective/builders/sm100_builder.inl
+270 -128
View File
@@ -582,7 +582,13 @@ sm100_sparse_get_tma_dispatch_policy() {
* Selected op also maximizes the TMEM_LOAD shape in order to minimize TMEM_LOADs issued,
* subject to the constraint of the provided per-warp tmem subpartition shape
**/
template<class GmemStrideTypeD, class ElementAccumulator, class ElementD, class TmemShape_MN, class FusionOp>
template<
class GmemStrideTypeD,
class ElementAccumulator,
class ElementD,
class TmemShape_MN,
bool IsBlockScaleSupported
>
constexpr auto
sm100_get_tmem_load_op() {
using namespace cute;
@@ -958,6 +964,172 @@ struct CallbacksBuilder<
>;
};
// Attempts to compute a reasonably performant epilogue tile or allows the user to provide one.
template<
class OpClass,
class CtaTileShape_MNK,
class EpilogueTileType,
class TmemWarpShape_MN,
class ElementC_,
class GmemStrideTypeC,
class ElementD,
class GmemStrideTypeD,
bool IsPerColScaleSupported
>
static constexpr auto
sm100_dense_compute_tile_shape_or_override() {
using namespace cute;
static_assert(!cute::is_same_v<OpClass, arch::OpClassSparseTensorOp> && !cute::is_same_v<OpClass, arch::OpClassBlockScaledSparseTensorOp>);
constexpr bool DisableSource = cute::is_void_v<ElementC_>;
using ElementC = cute::conditional_t<DisableSource, ElementD, ElementC_>;
if constexpr (is_same_v<OpClass, arch::OpClassBlockScaledTensorOp> &&
is_same_v<EpilogueTileType, EpilogueTileAuto> &&
size<1>(CtaTileShape_MNK{}) == 256) {
constexpr int CtaM = size<0>(CtaTileShape_MNK{});
constexpr int WarpM = size<0>(TmemWarpShape_MN{});
constexpr int DpFull = 32;
constexpr int M = cute::min(CtaM, DpFull * WarpM); // target 32dp tmem load
// Note:
// Set Epi_Tile_N to 128 support OverlappingAccum for the largest tile.
// This is a general workable epi_tile_N which does not promise best perf.
return make_tile(Int<M>{}, Int<128>{});
}
else if constexpr (is_same_v<EpilogueTileType, EpilogueTileAuto>) {
constexpr int CtaM = size<0>(CtaTileShape_MNK{});
constexpr int CtaN = size<1>(CtaTileShape_MNK{});
constexpr int WarpM = size<0>(TmemWarpShape_MN{});
constexpr int WarpN = size<1>(TmemWarpShape_MN{});
constexpr int MaxBits = cute::max(sizeof_bits_v<ElementC>, sizeof_bits_v<ElementD>);
constexpr int DpFull = 32; // tmem datapaths in 1 subpartition
constexpr int M = cute::min(CtaM, DpFull * WarpM); // target 32dp tmem load
constexpr int N_perf = [&]() constexpr { // Known subtile sizes tested for perf
// Epilogues w/o residual load are less sensitive to smem allocation
// Target a fixed amount of compute per epilogue iteration
if (DisableSource) {
if (MaxBits == 4) {
// Make epilogue tile larger to reduce the epilogue iterations.
// 64 is the experimental value. It will minimize epilogue iterations but keep the number of A/B buffers the same.
constexpr int ComputeElts = 8192;
return ComputeElts / M;
}
constexpr int ComputeElts = 4096;
return ComputeElts / M;
}
// Epilogues w/ residual load are more sensitive to smem allocation
// Target optimal smem distribution between epilogue+mainloop based on datatype+tilesize
else {
if (MaxBits == 32) {
return (CtaM > 64 && CtaN <= 128) ? 16 : 32;
}
// Per-column scaling is high register pressure, reduce tile to prevent spills
else if (IsPerColScaleSupported) {
return 32;
}
else if (MaxBits == 16) {
return (CtaN <= 128) ? 32 : 64;
}
else {
return 64;
}
}
}();
constexpr int N_min_C = (DisableSource || detail::is_m_major<GmemStrideTypeC>()) ? 8 * WarpN
: (sizeof_bits_v<ElementC> == 6) ? 128 * WarpN // TMA store only supports SW128B for FP6 data type
: 128 / sizeof_bits_v<ElementC> * WarpN;
constexpr int N_min_D = (detail::is_m_major<GmemStrideTypeD>()) ? 8 * WarpN
: (sizeof_bits_v<ElementD> == 6) ? 128 * WarpN // TMA store only supports SW128B for FP6 data type
: 128 / sizeof_bits_v<ElementD> * WarpN;
constexpr int N = cute::min(CtaN, cute::max(N_perf, N_min_C, N_min_D));
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
auto tile_m = Layout<Int<M>>{};
auto tile_n = Layout<Shape <Int<N / WarpN>,Int< WarpN>>,
Stride<Int< 1>,Int<CtaN / WarpN>>>{};
return make_tile(tile_m, coalesce(tile_n));
}
else {
static_assert(cute::is_tuple<EpilogueTileType>::value && not is_layout<EpilogueTileType>::value,
"EpilogueTile must be a cute::Tile or cute::Shape");
EpilogueTileType epi_tile;
constexpr int M = size<0>(shape(epi_tile));
constexpr int N = size<1>(shape(epi_tile));
static_assert(N % 8 == 0, "Unsupported tile shape");
return epi_tile;
}
}
template<
bool Is2SmMma,
class MmaTileShape_MNK
>
static constexpr auto
sm100_tmem_warps() {
if constexpr (Is2SmMma && size<0>(MmaTileShape_MNK{}) == 128) {
return Shape<_2,_2>{};
}
else {
return Shape<_4,_1>{};
}
}
template<
bool Is2SmMma,
class MmaTileShape_MNK
>
static constexpr auto
sm100_cta_tile_shape() {
if constexpr (Is2SmMma) { // 2x1 threadblock shape
auto [mma_tile_m, mma_tile_n, mma_tile_k] = MmaTileShape_MNK{};
auto cta_tile_m = reverse(shape_div(reverse(mma_tile_m), _2{})); // first MmaTile_M/2 elements, preserve multimode
return make_shape(cta_tile_m, mma_tile_n, mma_tile_k);
}
else { // 1x1 threadblock shape
return MmaTileShape_MNK{};
}
}
template<
class EpilogueScheduleType,
class ElementC_,
class ElementD,
int EpiTiles,
int FragmentSize
>
static constexpr auto
sm100_dense_dispatch_policy() {
// 8b residuals load fast and consume little smem, so the perf cost of waiting on stores to finish outweighs the cost of extra allocation
constexpr bool ReuseSmem = sizeof_bits_v<ElementC_> > 8;
// TMA store delay performs worse with residual loads
constexpr bool DelayTmaStore = is_void_v<ElementC_>;
constexpr int StagesD = cute::min(EpiTiles, 2);
constexpr int StagesC = ReuseSmem ? cute::max(cute::min(EpiTiles, 4), StagesD+1)
: cute::min(EpiTiles, 4);
if constexpr (is_base_of_v<PtrArrayNoSmemWarpSpecialized1Sm, EpilogueScheduleType> ||
is_base_of_v<PtrArrayNoSmemWarpSpecialized2Sm, EpilogueScheduleType>) {
return Sm100PtrArrayNoSmemWarpSpecialized{};
}
else if constexpr (is_base_of_v<NoSmemWarpSpecialized1Sm, EpilogueScheduleType> || is_base_of_v<NoSmemWarpSpecialized2Sm, EpilogueScheduleType>) {
return Sm100NoSmemWarpSpecialized{};
}
else if constexpr (is_same_v<EpilogueScheduleType, PtrArrayTmaWarpSpecialized1Sm> ||
is_same_v<EpilogueScheduleType, PtrArrayTmaWarpSpecialized2Sm>) {
constexpr bool DelayTmaStore_ = false; // TMA store delay complicates tensormap updates for Ptr-Array GEMMs
return Sm100PtrArrayTmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmem, DelayTmaStore_>{};
}
else {
return Sm100TmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmem, DelayTmaStore>{};
}
}
// Helper for building TMA warp-specialized collective epilogues, specialized by
// the fusion operation performed and the dispatch policy to use.
template <
@@ -1017,17 +1189,7 @@ private:
}
}
using CtaTileShape_MNK = decltype(cta_tile_shape());
static constexpr auto
tmem_warps() {
if constexpr (Is2SmMma && size<0>(MmaTileShape_MNK{}) == 128) {
return Shape<_2,_2>{};
}
else {
return Shape<_4,_1>{};
}
}
using TmemWarpShape_MN = decltype(tmem_warps());
using TmemWarpShape_MN = decltype(detail::sm100_tmem_warps<Is2SmMma, MmaTileShape_MNK>());
// Attempts to compute a reasonably performant epilogue tile or allows the user to provide one.
static constexpr auto
@@ -1041,84 +1203,10 @@ private:
ElementC_, GmemStrideTypeC, ElementD, GmemStrideTypeD, Schedule,
FusionOp>();
}
else if constexpr (is_same_v<OpClass, arch::OpClassBlockScaledTensorOp> &&
is_same_v<EpilogueTileType, EpilogueTileAuto> &&
size<1>(CtaTileShape_MNK{}) == 256) {
constexpr int CtaM = size<0>(CtaTileShape_MNK{});
constexpr int WarpM = size<0>(TmemWarpShape_MN{});
constexpr int DpFull = 32;
constexpr int M = cute::min(CtaM, DpFull * WarpM); // target 32dp tmem load
// Note:
// Set Epi_Tile_N to 128 support OverlappingAccum for the largest tile.
// This is a general workable epi_tile_N which does not promise best perf.
return make_tile(Int<M>{}, Int<128>{});
}
else if constexpr (is_same_v<EpilogueTileType, EpilogueTileAuto>) {
constexpr int CtaM = size<0>(CtaTileShape_MNK{});
constexpr int CtaN = size<1>(CtaTileShape_MNK{});
constexpr int WarpM = size<0>(TmemWarpShape_MN{});
constexpr int WarpN = size<1>(TmemWarpShape_MN{});
constexpr int MaxBits = cute::max(sizeof_bits_v<ElementC>, sizeof_bits_v<ElementD>);
constexpr int DpFull = 32; // tmem datapaths in 1 subpartition
constexpr int M = cute::min(CtaM, DpFull * WarpM); // target 32dp tmem load
constexpr int N_perf = [&]() constexpr { // Known subtile sizes tested for perf
// Epilogues w/o residual load are less sensitive to smem allocation
// Target a fixed amount of compute per epilogue iteration
if (DisableSource) {
if (MaxBits == 4) {
// Make epilogue tile larger to reduce the epilogue iterations.
// 64 is the experimental value. It will minimize epilogue iterations but keep the number of A/B buffers the same.
constexpr int ComputeElts = 8192;
return ComputeElts / M;
}
constexpr int ComputeElts = 4096;
return ComputeElts / M;
}
// Epilogues w/ residual load are more sensitive to smem allocation
// Target optimal smem distribution between epilogue+mainloop based on datatype+tilesize
else {
if (MaxBits == 32) {
return (CtaM > 64 && CtaN <= 128) ? 16 : 32;
}
// Per-column scaling is high register pressure, reduce tile to prevent spills
else if (FusionOp::IsPerColScaleSupported) {
return 32;
}
else if (MaxBits == 16) {
return (CtaN <= 128) ? 32 : 64;
}
else {
return 64;
}
}
}();
constexpr int N_min_C = (DisableSource || detail::is_m_major<GmemStrideTypeC>()) ? 8 * WarpN
: (sizeof_bits_v<ElementC> == 6) ? 128 * WarpN // TMA store only supports SW128B for FP6 data type
: 128 / sizeof_bits_v<ElementC> * WarpN;
constexpr int N_min_D = (detail::is_m_major<GmemStrideTypeD>()) ? 8 * WarpN
: (sizeof_bits_v<ElementD> == 6) ? 128 * WarpN // TMA store only supports SW128B for FP6 data type
: 128 / sizeof_bits_v<ElementD> * WarpN;
constexpr int N = cute::min(CtaN, cute::max(N_perf, N_min_C, N_min_D));
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
auto tile_m = Layout<Int<M>>{};
auto tile_n = Layout<Shape <Int<N / WarpN>,Int< WarpN>>,
Stride<Int< 1>,Int<CtaN / WarpN>>>{};
return make_tile(tile_m, coalesce(tile_n));
}
else {
static_assert(cute::is_tuple<EpilogueTileType>::value && not is_layout<EpilogueTileType>::value,
"EpilogueTile must be a cute::Tile or cute::Shape");
EpilogueTileType epi_tile;
constexpr int M = size<0>(shape(epi_tile));
constexpr int N = size<1>(shape(epi_tile));
static_assert(N % 8 == 0, "Unsupported tile shape");
return epi_tile;
return sm100_dense_compute_tile_shape_or_override<
OpClass, CtaTileShape_MNK, EpilogueTileType, TmemWarpShape_MN,
ElementC_, GmemStrideTypeC, ElementD, GmemStrideTypeD, FusionOp::IsPerColScaleSupported>();
}
}
using EpilogueTile_MN = decltype(epilogue_tile());
@@ -1129,30 +1217,18 @@ private:
using EpilogueWarpTileShape_MN = decltype(shape_div(EpilogueTileShape_MN{}, TmemWarpShape_MN{}));
using AccLoadOp = decltype(detail::sm100_get_tmem_load_op<
GmemStrideTypeD, ElementAccumulator, ElementD, EpilogueWarpTileShape_MN, FusionOp>());
GmemStrideTypeD, ElementAccumulator, ElementD, EpilogueWarpTileShape_MN,
FusionOp::IsBlockScaleSupported
>());
static constexpr auto
dispatch_policy() {
// 8b residuals load fast and consume little smem, so the perf cost of waiting on stores to finish outweighs the cost of extra allocation
constexpr bool ReuseSmem = sizeof_bits_v<ElementC_> > 8;
// TMA store delay performs worse with residual loads
constexpr bool DelayTmaStore = is_void_v<ElementC_>;
constexpr int StagesD = cute::min(EpiTiles, 2);
constexpr int StagesC = ReuseSmem ? cute::max(cute::min(EpiTiles, 4), StagesD+1)
: cute::min(EpiTiles, 4);
if constexpr (is_same_v<OpClass, arch::OpClassSparseTensorOp> ||
is_same_v<OpClass, arch::OpClassBlockScaledSparseTensorOp>) {
return detail::sparse::sm100_sparse_get_tma_dispatch_policy<CtaTileShape_MNK, EpilogueTile_MN, ElementC_, ElementD, Schedule>();
}
else if constexpr (is_same_v<Schedule, PtrArrayTmaWarpSpecialized1Sm> ||
is_same_v<Schedule, PtrArrayTmaWarpSpecialized2Sm>) {
constexpr bool DelayTmaStore_ = false; // TMA store delay complicates tensormap updates for Ptr-Array GEMMs
return Sm100PtrArrayTmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmem, DelayTmaStore_>{};
}
else {
return Sm100TmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmem, DelayTmaStore>{};
return detail::sm100_dense_dispatch_policy<Schedule, ElementC_, ElementD, EpiTiles, FragmentSize>();
}
}
@@ -1228,6 +1304,87 @@ public:
>;
};
template<
class OpClass,
class MmaTileShape_MNK,
class EpilogueTileType,
class ElementAccumulator_,
class ElementC,
class ElementD,
class Schedule,
class GmemStrideTypeC,
class GmemStrideTypeD,
bool IsPerColScaleSupported,
bool IsBlockScaleSupported
>
struct Sm100EpilogueDescriptor {
using ElementAccumulator = ElementAccumulator_;
static constexpr bool Is2SmMma = is_base_of_v<TmaWarpSpecialized2Sm, Schedule> || is_base_of_v<NoSmemWarpSpecialized2Sm, Schedule>;
using CtaTileShape_MNK = decltype(sm100_cta_tile_shape<Is2SmMma, MmaTileShape_MNK>());
using TileShape = CtaTileShape_MNK;
using TmemWarpShape_MN = decltype(detail::sm100_tmem_warps<Is2SmMma, MmaTileShape_MNK>());
using EpilogueTile = decltype(
sm100_dense_compute_tile_shape_or_override<OpClass, CtaTileShape_MNK, EpilogueTileType,
TmemWarpShape_MN, ElementC, GmemStrideTypeC, ElementD, GmemStrideTypeD, IsPerColScaleSupported>()
);
using EpilogueTileShape_MN = decltype(product_each(shape(EpilogueTile{})));
static constexpr int EpiTiles = size(shape_div(take<0,2>(CtaTileShape_MNK{}), EpilogueTileShape_MN{}));
static constexpr int FragmentSize = size(EpilogueTileShape_MN{}) / NumThreadsPerWarpGroup;
using DispatchPolicy = decltype(sm100_dense_dispatch_policy<Schedule, ElementC, ElementD, EpiTiles, FragmentSize>());
constexpr static int StagesC = DispatchPolicy::StagesC;
constexpr static int StagesD = DispatchPolicy::StagesD;
using EpilogueWarpTileShape_MN = decltype(shape_div(EpilogueTileShape_MN{}, TmemWarpShape_MN{}));
using AccLoadOp = decltype(detail::sm100_get_tmem_load_op<
GmemStrideTypeD, ElementAccumulator, ElementD, EpilogueWarpTileShape_MN,
IsBlockScaleSupported
>());
};
// Get Stride, SmemLayout, and CopyOpS2R for AuxLoad node
template<
typename EpilogueDescriptor,
typename StrideOrLayoutTag,
typename ElementAux
>
struct Sm100AuxLoadDescriptor {
constexpr static int Stages = EpilogueDescriptor::StagesC;
using EpilogueTile = typename EpilogueDescriptor::EpilogueTile;
using Element = ElementAux;
using Stride = cutlass::detail::TagToStrideC_t<StrideOrLayoutTag>;
using SmemLayoutAtom = decltype(detail::sm100_get_epilogue_smem_swizzle_layout_atom<
Stride, ElementAux, EpilogueTile>());
using CopyOpS2R = decltype(detail::sm100_get_smem_load_op<
Stride, ElementAux, typename EpilogueDescriptor::ElementAccumulator, typename EpilogueDescriptor::AccLoadOp>());
};
// Get Stride, SmemLayout, and CopyOpS2R for AuxStore node
template<
typename EpilogueDescriptor,
typename StrideOrLayoutTag,
typename ElementAux
>
struct Sm100AuxStoreDescriptor {
constexpr static int Stages = EpilogueDescriptor::StagesD;
using EpilogueTile = typename EpilogueDescriptor::EpilogueTile;
using Element = ElementAux;
using Stride = cutlass::detail::TagToStrideC_t<StrideOrLayoutTag>;
using SmemLayoutAtom = decltype(detail::sm100_get_epilogue_smem_swizzle_layout_atom<
Stride, ElementAux, EpilogueTile>());
using CopyOpR2S = decltype(detail::sm100_get_smem_store_op<
Stride, ElementAux, typename EpilogueDescriptor::ElementAccumulator, typename EpilogueDescriptor::AccLoadOp>());
};
} // namespace detail
///////////////////////////////////////////////////////////////////////////////
@@ -1304,17 +1461,7 @@ private:
}
}
using CtaTileShape_MNK = decltype(cta_tile_shape());
static constexpr auto
tmem_warps() {
if constexpr (Is2SmMma && size<0>(MmaTileShape_MNK{}) == 128) {
return Shape<_2,_2>{};
}
else {
return Shape<_4,_1>{};
}
}
using TmemWarpShape_MN = decltype(tmem_warps());
using TmemWarpShape_MN = decltype(detail::sm100_tmem_warps<Is2SmMma, MmaTileShape_MNK>());
static constexpr auto
epilogue_tile() {
@@ -1338,20 +1485,15 @@ private:
using EpilogueWarpTileShape_MN = decltype(shape_div(EpilogueTile{}, TmemWarpShape_MN{}));
using AccLoadOp = decltype(detail::sm100_get_tmem_load_op<
GmemStrideTypeD, ElementAccumulator, ElementD, EpilogueWarpTileShape_MN, FusionOp>());
GmemStrideTypeD, ElementAccumulator, ElementD, EpilogueWarpTileShape_MN,
FusionOp::IsBlockScaleSupported
>());
static constexpr int FragmentSize = size(EpilogueTile{}) / NumThreadsPerWarpGroup;
static constexpr auto
dispatch_policy() {
if constexpr (std::is_base_of_v<PtrArrayNoSmemWarpSpecialized1Sm, EpilogueScheduleType> ||
std::is_base_of_v<PtrArrayNoSmemWarpSpecialized2Sm, EpilogueScheduleType>) {
return Sm100PtrArrayNoSmemWarpSpecialized{};
}
else {
return Sm100NoSmemWarpSpecialized{};
}
}
using DispatchPolicy = decltype(dispatch_policy());
using EpilogueTileShape_MN = decltype(product_each(shape(EpilogueTile{})));
static constexpr int EpiTiles = size(shape_div(take<0,2>(CtaTileShape_MNK{}), EpilogueTileShape_MN{}));
using DispatchPolicy = decltype(detail::sm100_dense_dispatch_policy<EpilogueScheduleType, ElementC_, ElementD, EpiTiles, FragmentSize>());
static constexpr auto
fusion_callbacks() {
include/cutlass/epilogue/collective/sm90_epilogue_array_tma_warpspecialized.hpp
+1 -8
View File
@@ -507,8 +507,7 @@ public:
int thread_idx,
TensorStorage& shared_tensors,
TensorMapC const& load_tensormap,
int subtile_idx=-1,
bool wait_until_load_finishes = false) {
int subtile_idx=-1) {
using namespace cute;
// Indexing variables
@@ -595,12 +594,6 @@ public:
// Post-loop fusion callback entry point
pld_callbacks.end();
if (wait_until_load_finishes && did_load) {
typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_tma_consumer_state =
{last_load_producer_state.index(), !last_load_producer_state.phase(), last_load_producer_state.count()};
load_pipeline.consumer_wait(epi_load_pipe_tma_consumer_state);
}
return load_pipe_producer_state;
}
include/cutlass/gemm/collective/builders/sm100_blockscaled_mixed_tma_cpasync_umma_builder.inl
+274
View File
@@ -0,0 +1,274 @@
/***************************************************************************************************
* Copyright (c) 2025 - 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/gemm/collective/builders/sm100_common.inl"
#include "cutlass/gemm/collective/collective_builder_decl.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
namespace detail {
// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count.
template <
int CapacityBytes,
class ElementA,
class ElementB,
class TileShapeMNK,
class TileShapeSFA,
class TileShapeSFB,
int stages
>
constexpr int
sm100_compute_stage_count_or_override_blockscaled_mixed_tma_cpasync(StageCount<stages> stage_count) {
return stages;
}
template <
int CapacityBytes,
class ElementA,
class ElementB,
class TileShapeMNK,
class TileShapeSFA,
class TileShapeSFB,
int carveout_bytes
>
constexpr int
sm100_compute_stage_count_or_override_blockscaled_mixed_tma_cpasync(StageCountAutoCarveout<carveout_bytes> stage_count) {
// For MXF8F6F4 MMA, ElementA/B will be passed in as uint8_t
// Each stage include (CollectiveMma::SharedStorage)
// 1. smem for A and smem for B (CollectiveMma::SharedStorage::TensorStorage)
// 2. one MainloopPipeline = PipelineTmaUmmaAsync (CollectiveMma::SharedStorage::SharedStorage)
// 3. smem for SFB and smem for SFB (CollectiveMma::SharedStorage::TensorStorage, independent of input size b.c. sizeof(sf) is fixed)
constexpr auto mainloop_pipeline_bytes =
sizeof(typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage) +
sizeof(typename cutlass::PipelineUmmaConsumerAsync<1>::SharedStorage);
constexpr auto a_bits = cute::sizeof_bits_v<ElementA>;
constexpr auto b_bits = cute::sizeof_bits_v<ElementB>;
constexpr auto stage_sfa_bytes = size(filter_zeros(TileShapeSFA{}));
constexpr auto stage_sfb_bytes = size(filter_zeros(TileShapeSFB{}));
constexpr int stage_bytes =
cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
static_cast<int>(mainloop_pipeline_bytes + stage_sfa_bytes + stage_sfb_bytes);
return (CapacityBytes - carveout_bytes) / stage_bytes;
}
} // namespace detail
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ElementPairA,
class GmemLayoutATag,
int AlignmentA,
class ElementPairB,
class GmemLayoutBTag,
int AlignmentB,
class ElementAccumulator,
class TileShape_MNK,
class ClusterShape_MNK,
class StageCountType,
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm100,
arch::OpClassBlockScaledTensorOp,
ElementPairA,
GmemLayoutATag,
AlignmentA,
ElementPairB,
GmemLayoutBTag,
AlignmentB,
ElementAccumulator,
TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
ClusterShape_MNK, // Static cluster shape (_1, _1, _1)
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<cute::is_same_v<KernelMixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100, BuilderScheduleTag> >
>
{
using ElementSFA = typename detail::blockscaled::blockscaled_type<BuilderScheduleTag, ElementPairA>::sf_type;
using ElementSFB = typename detail::blockscaled::blockscaled_type<BuilderScheduleTag, ElementPairB>::sf_type;
using ElementA = typename detail::blockscaled::blockscaled_type<BuilderScheduleTag, ElementPairA>::data_type;
using ElementB = typename detail::blockscaled::blockscaled_type<BuilderScheduleTag, ElementPairB>::data_type;
using ElementSF = ElementSFA;
static constexpr cute::UMMA::Major UmmaMajorA = cutlass::gemm::collective::detail::tag_to_umma_major_A<GmemLayoutATag>();
static constexpr cute::UMMA::Major UmmaMajorB = cutlass::gemm::collective::detail::tag_to_umma_major_B<GmemLayoutBTag>();
static_assert(cute::is_static_v<TileShape_MNK>, "TileShape has to be static");
static_assert(detail::blockscaled::check_input_datatypes<BuilderScheduleTag, ElementPairA, ElementPairB, UmmaMajorA, UmmaMajorB>(), "Incorrect input types");
static constexpr bool is_2sm = false; // detail::blockscaled::is_2sm<TileShape_MNK, ClusterShape_MNK, BuilderScheduleTag>();
static constexpr auto Instr = detail::blockscaled::select_instr<ElementPairA, ElementPairB, ElementAccumulator, UmmaMajorA, UmmaMajorB, BuilderScheduleTag>();
using TiledMma = typename cutlass::gemm::collective::detail::TrivialBlockscaledMma<ElementPairA, ElementPairB, ElementAccumulator,
TileShape_MNK, ClusterShape_MNK,
UmmaMajorA, UmmaMajorB, Instr, BuilderScheduleTag, is_2sm>::type;
static constexpr bool UseMxf8f6f4 = Instr == detail::blockscaled::BlockScaledInstr::MXF4F6F8;
static_assert(UseMxf8f6f4 || (cutlass::gemm::detail::is_k_major_A<GmemLayoutATag>() && cutlass::gemm::detail::is_k_major_B<GmemLayoutBTag>()), "Only MMA.MXF8F6F4 supports non-K major inputs");
// Data type used by MMA instruction
using ElementAMma = decltype(cutlass::gemm::collective::detail::sm1xx_kernel_input_element_to_mma_input_element<ElementA, UseMxf8f6f4>());
using ElementBMma = decltype(cutlass::gemm::collective::detail::sm1xx_kernel_input_element_to_mma_input_element<ElementB, UseMxf8f6f4>());
static_assert(detail::sm1xx_gemm_check_for_f8f6f4_mix8bit_requirement<ElementAMma, ElementBMma,
TileShape_MNK, ClusterShape_MNK,
GmemLayoutATag, GmemLayoutBTag, false /*is_sparse*/, is_2sm>(),
"TileSize and MNK Major does not met with MMA Mix 8-bit TMA load requirement" );
static constexpr uint32_t SFVectorSize = TiledMma::SFVecSize;
using ElementAMma_SmemAllocType = cute::conditional_t<UseMxf8f6f4, uint8_t, ElementAMma>;
using ElementBMma_SmemAllocType = cute::conditional_t<UseMxf8f6f4, uint8_t, ElementBMma>;
// using TiledMma = decltype(detail::sm100_make_trivial_tiled_mma<
// ElementAMma, ElementBMma, ElementAccumulator,
// decltype(cute::product_each(TileShape_MNK{})), ClusterShape_MNK,
// UmmaMajorA, UmmaMajorB, BuilderScheduleTag>());
using AtomThrID = typename TiledMma::AtomThrID;
using Sm1xxBlkScaledConfig = cutlass::detail::Sm1xxBlockScaledConfig<SFVectorSize>;
// ((MMA_TILE_M,MMA_TILE_K), MMA_M, MMA_K)
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(cute::size<0>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
// ((MMA_TILE_N,MMA_TILE_K), MMA_N, MMA_K)
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(cute::size<1>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
// Assigning 4 warps for mainloop load of B
static constexpr int NumLoadThreadsCpAsync = 128;
using SmemShapeA_M = decltype(shape_div(shape<0>(TileShape_MNK{}), shape_div(shape<0>(TileShape_MNK{}), size<0>(TileShape_MNK{}) / size(AtomThrID{}))));
using SmemShapeA_K = decltype(cute::get<2>(TileShape_MNK{}));
using GmemTiledCopyA = decltype(cutlass::gemm::collective::detail::sm100_cluster_shape_to_tma_atom_A(
ClusterShape_MNK{}, AtomThrID{}));
using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
UmmaMajorA, ElementAMma_SmemAllocType, SmemShapeA_M, SmemShapeA_K>());
using AlignmentTypeB = cute::uint_byte_t<static_cast<int>(cutlass::sizeof_bits<ElementB>::value) * AlignmentB / 8>;
using GmemCopyAtomB = cute::Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL_ZFILL<AlignmentTypeB>, ElementB>;
using GmemTiledCopyB = decltype(detail::make_simt_gmem_tiled_copy<
GmemCopyAtomB, NumLoadThreadsCpAsync, AlignmentB, TagToStrideB_t<GmemLayoutBTag>,
decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
using BlockTileB_N = decltype(cute::size<0,0>(MmaShapeB_NK{}) * cute::size<1>(MmaShapeB_NK{}));
using BlockTileB_K = decltype(cute::size<0,1>(MmaShapeB_NK{}) * cute::size<2>(MmaShapeB_NK{}));
using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
UmmaMajorB, ElementBMma_SmemAllocType, BlockTileB_N, BlockTileB_K>());
using GmemTiledCopySFA = decltype(cutlass::gemm::collective::detail::sm100_cluster_shape_to_tma_atom_A(
ClusterShape_MNK{}, AtomThrID{}));
using GmemTiledCopySFB = decltype(cutlass::gemm::collective::detail::sm100_cluster_shape_to_tma_atom_SFB(
ClusterShape_MNK{}, AtomThrID{}));
using GmemTiledCopyPairA = decltype(cute::make_tuple(GmemTiledCopyA{}, GmemTiledCopySFA{}));
using GmemTiledCopyPairB = decltype(cute::make_tuple(GmemTiledCopyB{}, GmemTiledCopySFB{}));
using SmemLayoutAtomSFA = decltype(Sm1xxBlkScaledConfig::deduce_smem_layoutSFA(TiledMma{}, TileShape_MNK{}));
using SmemLayoutAtomSFB = decltype(Sm1xxBlkScaledConfig::deduce_smem_layoutSFB(TiledMma{}, TileShape_MNK{}));
using SmemLayoutAtomsA = decltype(cute::make_tuple(SmemLayoutAtomA{}, SmemLayoutAtomSFA{}));
using SmemLayoutAtomsB = decltype(cute::make_tuple(SmemLayoutAtomB{}, SmemLayoutAtomSFB{}));
using StrideA = cutlass::gemm::TagToStrideA_t<GmemLayoutATag>;
using StrideB = cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>;
using InternalStrideA = cute::remove_pointer_t<StrideA>;
using InternalStrideB = cute::remove_pointer_t<StrideB>;
using InternalLayoutSFA = decltype(Sm1xxBlkScaledConfig::deduce_layoutSFA());
using InternalLayoutSFB = decltype(Sm1xxBlkScaledConfig::deduce_layoutSFB());
using LayoutSFA = cute::conditional_t<cute::is_same_v<InternalStrideA, StrideA>, InternalLayoutSFA, InternalLayoutSFA *>;
using LayoutSFB = cute::conditional_t<cute::is_same_v<InternalStrideB, StrideB>, InternalLayoutSFB, InternalLayoutSFB *>;
using StridePairA = decltype(cute::make_tuple(StrideA{}, LayoutSFA{}));
using StridePairB = decltype(cute::make_tuple(StrideB{}, LayoutSFB{}));
static constexpr uint32_t AccumulatorPipelineStageCount = 2;
// Calculate scheduler pipeline stages. Having one more stage than the accumulator allows more latency hiding.
static constexpr uint32_t SchedulerPipelineStageCount = AccumulatorPipelineStageCount + 1;
// AccumulatorPipeline = PipelineUmmaAsync
static constexpr auto AccumulatorPipelineStorage = sizeof(typename cutlass::PipelineUmmaAsync<AccumulatorPipelineStageCount>::SharedStorage);
// CLCPipeline = PipelineCLCFetchAsync
static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape_MNK>::SharedStorage);
// CLC (scheduler) response
static constexpr auto CLCResponseStorage = SchedulerPipelineStageCount * detail::CLCResponseSize;
// Smem usage that's not part of CollectiveEpilogue::SharedStorage & CollectiveMainloop::SharedStorage
static constexpr auto KernelSmemCarveout = static_cast<int>( AccumulatorPipelineStorage +
CLCPipelineStorage +
CLCResponseStorage);
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int Sm100ReducedSmemCapacityBytes = cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
using SmemTileShape = cute::Shape<SmemShapeA_M, BlockTileB_N, SmemShapeA_K>;
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_blockscaled_mixed_tma_cpasync<
Sm100ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, SmemLayoutAtomSFA, SmemLayoutAtomSFB>(StageCountType{});
static_assert(PipelineStages > 0, "Smem usage is too high. Can't create any SMEM buffers for A, B, SFA, and SFB.");
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
cutlass::gemm::MainloopSm100UmmaMixedTmaCpAsyncWarpSpecializedBlockScaled<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>,
TileShape_MNK,
cute::tuple<ElementA, ElementSF>,
StridePairA,
cute::tuple<ElementB, ElementSF>,
StridePairB,
TiledMma,
GmemTiledCopyPairA,
SmemLayoutAtomsA,
void,
cute::identity,
GmemTiledCopyPairB,
SmemLayoutAtomsB,
void,
cute::identity
>;
};
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl
+1
View File
@@ -120,6 +120,7 @@ struct CollectiveBuilder<
BuilderScheduleTag,
cute::enable_if_t<
// Blockscaled Gemm
(not cute::is_same_v<KernelMixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100, BuilderScheduleTag>) &&
(cute::is_base_of_v<KernelScheduleBlockScaledGemmSm100, BuilderScheduleTag> ||
cute::is_same_v<KernelScheduleAuto, BuilderScheduleTag>)
&&
include/cutlass/gemm/collective/builders/sm100_mixed_tma_cpasync_umma_builder.inl
+171
View File
@@ -0,0 +1,171 @@
/***************************************************************************************************
* Copyright (c) 2025 - 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/gemm/collective/builders/sm100_common.inl"
#include "cutlass/gemm/collective/collective_builder_decl.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ElementA,
class GmemLayoutATag,
int AlignmentA,
class ElementB,
class GmemLayoutBTag,
int AlignmentB,
class ElementAccumulator,
class TileShape_MNK,
class ClusterShape_MNK,
class StageCountType,
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm100,
arch::OpClassTensorOp,
ElementA,
GmemLayoutATag,
AlignmentA,
ElementB,
GmemLayoutBTag,
AlignmentB,
ElementAccumulator,
TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
ClusterShape_MNK, // Static cluster shape (_1, _1, _1)
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<cute::is_same_v<KernelMixedTmaCpAsyncWarpSpecialized1SmSm100, BuilderScheduleTag> >
>
{
static_assert(cute::is_static_v<TileShape_MNK>, "TileShape has to be static");
static constexpr cute::UMMA::Major UmmaMajorA = cutlass::gemm::collective::detail::tag_to_umma_major_A<GmemLayoutATag>();
static constexpr cute::UMMA::Major UmmaMajorB = cutlass::gemm::collective::detail::tag_to_umma_major_B<GmemLayoutBTag>();
// Data type used by MMA instruction
using ElementAMma = decltype(cutlass::gemm::collective::detail::sm1xx_kernel_input_element_to_mma_input_element<ElementA>());
using ElementBMma = decltype(cutlass::gemm::collective::detail::sm1xx_kernel_input_element_to_mma_input_element<ElementB>());
using ElementAMma_SmemAllocType = cute::conditional_t<cute::sizeof_bits_v<ElementAMma> < 8, uint8_t, ElementAMma>;
using ElementBMma_SmemAllocType = cute::conditional_t<cute::sizeof_bits_v<ElementBMma> < 8, uint8_t, ElementBMma>;
using TiledMma = decltype(detail::sm100_make_trivial_tiled_mma<
ElementAMma, ElementBMma, ElementAccumulator,
decltype(cute::product_each(TileShape_MNK{})), ClusterShape_MNK,
UmmaMajorA, UmmaMajorB, BuilderScheduleTag>());
using AtomThrID = typename TiledMma::AtomThrID;
// ((MMA_TILE_M,MMA_TILE_K), MMA_M, MMA_K)
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(cute::size<0>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
// ((MMA_TILE_N,MMA_TILE_K), MMA_N, MMA_K)
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(cute::size<1>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
// Assigning 4 warps for mainloop load of B
static constexpr int NumLoadThreadsCpAsync = 128;
using SmemShapeA_M = decltype(shape_div(shape<0>(TileShape_MNK{}), shape_div(shape<0>(TileShape_MNK{}), size<0>(TileShape_MNK{}) / size(AtomThrID{}))));
using SmemShapeA_K = decltype(cute::get<2>(TileShape_MNK{}));
using GmemTiledCopyA = decltype(cutlass::gemm::collective::detail::sm100_cluster_shape_to_tma_atom_A(
ClusterShape_MNK{}, AtomThrID{}));
using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
UmmaMajorA, ElementAMma_SmemAllocType, SmemShapeA_M, SmemShapeA_K>());
using AlignmentTypeB = cute::uint_byte_t<static_cast<int>(sizeof(ElementB)) * AlignmentB>;
using GmemCopyAtomB = cute::Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL_ZFILL<AlignmentTypeB>, ElementB>;
using GmemTiledCopyB = decltype(detail::make_simt_gmem_tiled_copy<
GmemCopyAtomB, NumLoadThreadsCpAsync, AlignmentB, TagToStrideB_t<GmemLayoutBTag>,
decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
using BlockTileB_N = decltype(cute::size<0,0>(MmaShapeB_NK{}) * cute::size<1>(MmaShapeB_NK{}));
using BlockTileB_K = decltype(cute::size<0,1>(MmaShapeB_NK{}) * cute::size<2>(MmaShapeB_NK{}));
using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
UmmaMajorB, ElementBMma_SmemAllocType, BlockTileB_N, BlockTileB_K>());
static constexpr uint32_t AccumulatorPipelineStageCount = 2;
// Calculate scheduler pipeline stages. Having one more stage than the accumulator allows more latency hiding.
static constexpr uint32_t SchedulerPipelineStageCount = AccumulatorPipelineStageCount + 1;
// AccumulatorPipeline = PipelineUmmaAsync
static constexpr auto AccumulatorPipelineStorage = sizeof(typename cutlass::PipelineUmmaAsync<AccumulatorPipelineStageCount>::SharedStorage);
// CLCPipeline = PipelineCLCFetchAsync
static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape_MNK>::SharedStorage);
// CLC (scheduler) response
static constexpr auto CLCResponseStorage = SchedulerPipelineStageCount * detail::CLCResponseSize;
// Smem usage that's not part of CollectiveEpilogue::SharedStorage & CollectiveMainloop::SharedStorage
static constexpr auto KernelSmemCarveout = static_cast<int>( AccumulatorPipelineStorage +
CLCPipelineStorage +
CLCResponseStorage);
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int Sm100ReducedSmemCapacityBytes = cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
using SmemTileShape = cute::Shape<SmemShapeA_M, BlockTileB_N, SmemShapeA_K>;
using MainloopPipelineStorage = typename cutlass::PipelineUmmaConsumerAsync<1>::SharedStorage;
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override<
Sm100ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, MainloopPipelineStorage>(StageCountType{});
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
cutlass::gemm::MainloopSm100UmmaMixedTmaCpAsyncWarpSpecialized<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>,
TileShape_MNK,
ElementA,
cutlass::gemm::TagToStrideA_t<GmemLayoutATag>,
ElementB,
cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>,
TiledMma,
GmemTiledCopyA,
SmemLayoutAtomA,
void,
cute::identity,
GmemTiledCopyB,
SmemLayoutAtomB,
void,
cute::identity
>;
};
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/builders/sm100_umma_builder.inl
+1
View File
@@ -184,6 +184,7 @@ struct CollectiveBuilder<
not cute::is_complex_v<ElementA> && not cute::is_complex_v<ElementB> &&
// Dense Gemm / PtrArrayDenseGemm
(
(not cute::is_same_v<KernelMixedTmaCpAsyncWarpSpecialized1SmSm100, BuilderScheduleTag>) &&
(not cute::is_same_v<KernelWarpSpecialized1SmSm100, BuilderScheduleTag>) &&
(cute::is_base_of_v<KernelScheduleSm100DenseGemm, BuilderScheduleTag> ||
cute::is_same_v<KernelScheduleAuto, BuilderScheduleTag>)) &&
include/cutlass/gemm/collective/builders/sm1xx_common.inl
+3
View File
@@ -502,6 +502,7 @@ check_input_datatypes() {
|| (cute::is_same_v<BuilderScheduleTag, KernelScheduleBlockScaledGemmSm100>)
|| (cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecialized1SmBlockScaledSm100>)
|| (cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecialized2SmBlockScaledSm100>)
|| (cute::is_same_v<BuilderScheduleTag, KernelMixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100>)
// SM100 BS ptr_array
|| (cute::is_same_v<BuilderScheduleTag, KernelSchedulePtrArrayBlockScaledGemmSm100>)
|| (cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized1SmBlockScaledSm100>)
@@ -578,6 +579,8 @@ check_input_datatypes() {
((SfVectorSizeA == 32 && cute::is_same_v<KernelScheduleAuto, BuilderScheduleTag>)
|| (SfVectorSizeA == 32 && cute::is_same_v<KernelTmaWarpSpecializedPingpong, BuilderScheduleTag>)
|| (SfVectorSizeA == 32 && cute::is_same_v<KernelTmaWarpSpecializedCooperative, BuilderScheduleTag>)
|| (SfVectorSizeA == 32 && cute::is_same_v<KernelPtrArrayTmaWarpSpecializedPingpong, BuilderScheduleTag>)
|| (SfVectorSizeA == 32 && cute::is_same_v<KernelPtrArrayTmaWarpSpecializedCooperative, BuilderScheduleTag>)
|| (SfVectorSizeA == 32 && cute::is_base_of_v<KernelScheduleBlockScaledGemmSm100, BuilderScheduleTag>)
|| (SfVectorSizeA == 32 && cute::is_base_of_v<KernelSchedulePtrArrayBlockScaledGemmSm100, BuilderScheduleTag>)
|| (SfVectorSizeA == 64 && cute::is_base_of_v<KernelScheduleBlockScaledSparseGemmSm100, BuilderScheduleTag>)
include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl
+7 -7
View File
@@ -1069,10 +1069,10 @@ struct CollectiveBuilder<
StageCountType,
KernelScheduleType,
cute::enable_if_t<
(cute::is_same_v<KernelScheduleType, KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum> or
cute::is_same_v<KernelScheduleType, KernelPtrArrayTmaWarpSpecializedCooperativeFP8BlockScaledAccum> or
cute::is_same_v<KernelScheduleType, KernelTmaWarpSpecializedPingpongFP8BlockScaledAccum> or
cute::is_same_v<KernelScheduleType, KernelPtrArrayTmaWarpSpecializedPingpongFP8BlockScaledAccum>) and
(cute::is_same_v<KernelScheduleType, KernelTmaWarpSpecializedCooperativeFP8Blockwise> or
cute::is_same_v<KernelScheduleType, KernelPtrArrayTmaWarpSpecializedCooperativeFP8Blockwise> or
cute::is_same_v<KernelScheduleType, KernelTmaWarpSpecializedPingpongFP8Blockwise> or
cute::is_same_v<KernelScheduleType, KernelPtrArrayTmaWarpSpecializedPingpongFP8Blockwise>) and
not detail::is_use_rmem_A<ElementA, GmemLayoutPairA, ElementB, GmemLayoutPairB>()
>
> {
@@ -1105,7 +1105,7 @@ struct CollectiveBuilder<
cute::is_base_of_v<KernelPtrArrayTmaWarpSpecializedPingpong, KernelScheduleType>);
static constexpr bool IsFP8Input = detail::is_input_fp8<ElementA, ElementB>();
static_assert(IsFP8Input, "Warp Specialized gemm with FP8 BlockScaled Accumulator is only compatible with FP8 Blocked Scaled version right now.");
static_assert(IsFP8Input, "Warp Specialized gemm with FP8 Blockwise (Software) Scaling is only compatible with FP8 inputs version right now.");
// For fp32 types, map to tf32 MMA value type
using ElementAMma = cute::conditional_t<cute::is_same_v<ElementA, float>, tfloat32_t, ElementA>;
@@ -1146,8 +1146,8 @@ struct CollectiveBuilder<
static constexpr int PipelineStages = detail::compute_stage_count_with_blockwise_scale<detail::sm90_smem_capacity_bytes - KernelSmemCarveout,
ElementAMma, ElementBMma, ElementBlockScale, TileShape_MNK, ScaleMsPerTile, ScaleNsPerTile>(StageCountType{});
using DispatchPolicy = cute::conditional_t<IsArrayOfPointersGemm,
MainloopSm90ArrayTmaGmmaWarpSpecializedBlockScaling<PipelineStages, ClusterShape_MNK, KernelScheduleType>,
MainloopSm90TmaGmmaWarpSpecializedBlockScalingFP8<PipelineStages, ClusterShape_MNK, KernelScheduleType>>;
MainloopSm90ArrayTmaGmmaWarpSpecializedBlockwise<PipelineStages, ClusterShape_MNK, KernelScheduleType>,
MainloopSm90TmaGmmaWarpSpecializedBlockwiseFP8<PipelineStages, ClusterShape_MNK, KernelScheduleType>>;
using SmemCopyAtomA = void;
using SmemCopyAtomB = void;
include/cutlass/gemm/collective/collective_builder.hpp
+2
View File
@@ -49,6 +49,8 @@
#include "cutlass/gemm/collective/builders/sm100_simt_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_mixed_input_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_cpasync_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_mixed_tma_cpasync_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_blockscaled_mixed_tma_cpasync_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm103_blockscaled_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm120_mma_builder.inl"
#include "cutlass/gemm/collective/builders/sm120_blockscaled_mma_builder.inl"
include/cutlass/gemm/collective/collective_mma.hpp
+2
View File
@@ -65,6 +65,8 @@
#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_blockwise_scaling.hpp"
#include "cutlass/gemm/collective/sm100_mma_warpspecialized_mixed_input.hpp"
#include "cutlass/gemm/collective/sm100_mma_cpasync_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm100_mma_mixed_tma_cpasync_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm100_blockscaled_mma_mixed_tma_cpasync_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm103_blockscaled_mma_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm103_blockscaled_mma_array_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm120_mma_tma.hpp"
include/cutlass/gemm/collective/sm100_blockscaled_mma_mixed_tma_cpasync_warpspecialized.hpp
+1043
View File
File diff suppressed because it is too large Load Diff
include/cutlass/gemm/collective/sm100_mma_mixed_tma_cpasync_warpspecialized.hpp
+758
View File
@@ -0,0 +1,758 @@
/***************************************************************************************************
* Copyright (c) 2025 - 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"
#include "cutlass/detail/cluster.hpp"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/numeric_types.h"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/gemm/gemm.h"
#include "cutlass/trace.h"
#include "cutlass/kernel_hardware_info.hpp"
#include "cutlass/arch/memory.h"
#include "cute/algorithm/functional.hpp"
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
#include "cutlass/gemm/collective/collective_mma_decl.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
using namespace cute;
/////////////////////////////////////////////////////////////////////////////////////////////////
// WarpSpecialized Mainloop
// Both DMA Load and MMA methods of this class must be run by a single thread that's picked by elect_one
template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
class StrideA_,
class ElementB_,
class StrideB_,
class TiledMma_,
class GmemTiledCopyA_,
class SmemLayoutAtomA_,
class SmemCopyAtomA_,
class TransformA_,
class GmemTiledCopyB_,
class SmemLayoutAtomB_,
class SmemCopyAtomB_,
class TransformB_>
struct CollectiveMma<
MainloopSm100UmmaMixedTmaCpAsyncWarpSpecialized<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
TileShape_,
ElementA_,
StrideA_,
ElementB_,
StrideB_,
TiledMma_,
GmemTiledCopyA_,
SmemLayoutAtomA_,
SmemCopyAtomA_,
TransformA_,
GmemTiledCopyB_,
SmemLayoutAtomB_,
SmemCopyAtomB_,
TransformB_>
{
using TiledMma = TiledMma_;
using AtomThrShapeMNK = Shape<decltype(shape<0>(typename TiledMma::ThrLayoutVMNK{})), _1, _1>;
// Statically asserting to ensure only 1x1x1 cluster shape & 1sm setup is received
static_assert(size(AtomThrShapeMNK{}) == 1, "Lower alignment SM100 GEMM only supports 1SM MMA");
static_assert(size(ClusterShape{}) == 1, "CPASYNC does not support multicast so the cluster shape is restricted to 1, 1, 1");
static_assert(size(typename TiledMma::AtomThrID{}) == 1);
using DispatchPolicy = MainloopSm100UmmaMixedTmaCpAsyncWarpSpecialized<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
// TileShape refers to MmaTileShape to adapt for runtime cluster
using TileShape = TileShape_;
CUTE_STATIC_ASSERT_V(evenly_divides(TileShape{}, tile_shape(TiledMma{})),
"Static cluster shape used: TileShape should be evenly divided by TiledMma");
// Define A and B block shapes
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{}))));
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{}))));
// using LoadShapeA_MK = decltype(select<0,2>(TileShape{}));
using LoadShapeB_NK = decltype(select<1,2>(TileShape{}));
// CtaShape_MNK is queried from collective in all kernel layers
using CtaShape_MNK = TileShape;
using ElementA = ElementA_;
using ElementAMma = typename TiledMma::ValTypeA;
using StrideA = StrideA_;
using ElementB = ElementB_;
using ElementBMma = typename TiledMma::ValTypeB;
using StrideB = StrideB_;
static constexpr bool IsRuntimeDataTypeA = cute::is_same_v<ElementA, cutlass::type_erased_dynamic_float8_t>;
static constexpr bool IsRuntimeDataTypeB = cute::is_same_v<ElementB, cutlass::type_erased_dynamic_float8_t>;
static_assert(IsRuntimeDataTypeA == IsRuntimeDataTypeB,
"ElementA and ElementB should be both runtime or both static.");
static constexpr bool IsRuntimeDataType = IsRuntimeDataTypeA && IsRuntimeDataTypeB;
using ElementAccumulator = typename TiledMma::ValTypeC;
using GmemTiledCopyA = GmemTiledCopyA_;
using GmemTiledCopyB = GmemTiledCopyB_;
using SmemLayoutAtomA = SmemLayoutAtomA_;
using SmemLayoutAtomB = SmemLayoutAtomB_;
using SmemCopyAtomA = SmemCopyAtomA_;
using SmemCopyAtomB = SmemCopyAtomB_;
using TransformA = TransformA_;
using TransformB = TransformB_;
using ArchTag = typename DispatchPolicy::ArchTag;
using MainloopPipelineTMA = cutlass::PipelineTmaUmmaAsync<DispatchPolicy::Stages, ClusterShape, AtomThrShapeMNK>;
using MainloopPipelineTMAState = typename MainloopPipelineTMA::PipelineState;
using MainloopPipelineCpAsync = cutlass::PipelineUmmaConsumerAsync<DispatchPolicy::Stages, AtomThrShapeMNK>;
using MainloopPipelineCpAsyncState = typename MainloopPipelineCpAsync::PipelineState;
// static_assert(size(GmemTiledCopyA{}) == size(GmemTiledCopyB{}), "A and B GmemTiledCopy should share the same thread count");
static constexpr int NumLoadThreadsCpAsync = size(GmemTiledCopyB{});
static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtomA must be rank 2 (M,K)");
static_assert(((size<0,0>(MmaShapeA_MK{}) * size<1>(MmaShapeA_MK{})) % size<0>(SmemLayoutAtomA{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,1>(MmaShapeA_MK{}) * size<2>(MmaShapeA_MK{})) % size<1>(SmemLayoutAtomA{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(cute::is_void_v<SmemCopyAtomA>,
"SM100 UMMA cannot have a non-void copy atom for smem sourced instructions.");
static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtomB must be rank 2 (N,K)");
static_assert(((size<0,0>(MmaShapeB_NK{}) * size<1>(MmaShapeB_NK{})) % size<0>(SmemLayoutAtomB{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,1>(MmaShapeB_NK{}) * size<2>(MmaShapeB_NK{})) % size<1>(SmemLayoutAtomB{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(cute::is_void_v<SmemCopyAtomB>,
"SM100 UMMA cannot have a non-void copy atom for smem sourced instructions.");
// Tile along K mode first before tiling over MN. PIPE mode last as usual.
// (MMA_TILE_M,MMA_TILE_K),MMA_M,MMA_K,PIPE)
using SmemLayoutA = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomA{},
append(MmaShapeA_MK{}, Int<DispatchPolicy::Stages>{}),
conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
using MmaSmemLayoutB = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomB{},
append(MmaShapeB_NK{}, Int<DispatchPolicy::Stages>{}),
conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
using LoadSmemLayoutB = decltype(tile_to_shape(
SmemLayoutAtomB{},
append(LoadShapeB_NK{}, Int<DispatchPolicy::Stages>{}),
conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more.");
static_assert(cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
"MMA atom must source both A and B operand from smem_desc for this mainloop.");
using TmaInternalElementA = cute::conditional_t<cute::is_same_v<ElementA, float>, cutlass::tfloat32_t, ElementAMma>;
using SmemAllocTypeA = cute::conditional_t<cute::sizeof_bits_v<ElementAMma> < 8, uint8_t, ElementAMma>;
using SmemAllocTypeB = cute::conditional_t<cute::sizeof_bits_v<ElementBMma> < 8, uint8_t, ElementBMma>;
using BitTypeElementA = cute::uint_bit_t<cute::sizeof_bits_v<ElementA>>;
using BitTypeElementB = cute::uint_bit_t<cute::sizeof_bits_v<ElementB>>;
using ArrayElementA = cute::conditional_t<IsRuntimeDataTypeA, BitTypeElementA, ElementA>;
using ArrayElementB = cute::conditional_t<IsRuntimeDataTypeB, BitTypeElementB, ElementB>;
using RuntimeDataTypeA = cute::conditional_t<IsRuntimeDataTypeA, cute::UMMA::MXF8F6F4Format, void*>;
using RuntimeDataTypeB = cute::conditional_t<IsRuntimeDataTypeB, cute::UMMA::MXF8F6F4Format, void*>;
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128, _0> {
cute::array_aligned<SmemAllocTypeA, cute::cosize_v<SmemLayoutA>> smem_A;
cute::array_aligned<SmemAllocTypeB, cute::cosize_v<LoadSmemLayoutB>> smem_B;
} tensors;
using PipelineStorageTMA = typename MainloopPipelineTMA::SharedStorage;
using PipelineStorageCpAsync = typename MainloopPipelineCpAsync::SharedStorage;
struct PipelineStorage : cute::aligned_struct<16, _0> {
alignas(16) PipelineStorageTMA tma;
alignas(16) PipelineStorageCpAsync cpasync;
} pipelines;
};
// Expose shared storage for tensors/pipelines separately to allow kernel layer to reorder them.
using TensorStorage = typename SharedStorage::TensorStorage;
using PipelineStorage = typename SharedStorage::PipelineStorage;
static constexpr uint32_t TmaTransactionBytes =
cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * cosize(take<0,3>(SmemLayoutA{})) * cute::sizeof_bits_v<ElementA>);
template <class AccTensor>
struct TmemStorage {
AccTensor accumulators;
};
// 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{};
};
// Device side kernel params
struct Params {
using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(ClusterShape{}),
make_tile(typename TiledMma::AtomThrID{})));
using TMA_A = decltype(make_tma_atom_A_sm100<TmaInternalElementA>(
GmemTiledCopyA{},
make_tensor(recast_ptr<TmaInternalElementA>(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}),
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
ClusterLayout_VMNK{})
);
TMA_A tma_load_a;
ArrayElementB const* ptr_B{nullptr};
StrideB dB{};
RuntimeDataTypeA runtime_data_type_a;
RuntimeDataTypeB runtime_data_type_b;
};
CUTLASS_DEVICE
CollectiveMma(Params const& params)
: runtime_data_type_a_(params.runtime_data_type_a)
, runtime_data_type_b_(params.runtime_data_type_b) {
observed_tma_load_a_ = &params.tma_load_a;
}
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(
ProblemShape const& problem_shape,
Arguments const& args,
[[maybe_unused]] void* workspace,
cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) {
// Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
auto ptr_A = recast_ptr<TmaInternalElementA>(args.ptr_A);
auto ptr_B = recast_ptr<ElementBMma>(args.ptr_B);
Tensor tensor_a = make_tensor(ptr_A, make_layout(make_shape(M,K,L), args.dA));
auto cluster_layout_vmnk = tiled_divide(make_layout(ClusterShape{}), make_tile(typename TiledMma::AtomThrID{}));
typename Params::TMA_A tma_load_a = make_tma_atom_A_sm100<TmaInternalElementA>(
GmemTiledCopyA{},
tensor_a,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
return {
tma_load_a,
args.ptr_B,
args.dB,
args.runtime_data_type_a,
args.runtime_data_type_b
};
}
template <class ProblemShape>
static bool
can_implement(
ProblemShape const& problem_shape,
[[maybe_unused]] Arguments const& args) {
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
static constexpr bool IsF8F6F4 = detail::is_sm100_mma_f8f6f4<TiledMma, ElementA, ElementB>();
constexpr int tma_alignment_bits_A = cutlass::detail::get_input_alignment_bits<ElementA, IsF8F6F4>();
constexpr int min_tma_aligned_elements_A = tma_alignment_bits_A / cute::sizeof_bits<ElementA>::value;
bool implementable = true;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
}
implementable = implementable && cutlass::detail::check_alignment<GmemTiledCopyB::NumValSrc>(cute::make_shape(N,K,L), StrideB{});
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for CpAsync.\n");
}
return implementable;
}
/// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
CUTLASS_DEVICE void
prefetch_tma_descriptors() {
cute::prefetch_tma_descriptor(observed_tma_load_a_->get_tma_descriptor());
}
/// Construct A Single Stage's Accumulator Shape
CUTLASS_DEVICE static
auto
partition_accumulator_shape() {
auto acc_shape = partition_shape_C(TiledMma{}, take<0,2>(TileShape{})); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N)
return acc_shape;
}
template <class TmemStorage>
CUTLASS_DEVICE static
auto
slice_accumulator(TmemStorage tmem_storage, int stage) {
return cute::make_tuple(tmem_storage.accumulators(_,_,_,stage));
}
template <class EpilogueTile, bool IsOverlappingAccum = false>
CUTLASS_DEVICE static
auto
init_tmem_tensors(EpilogueTile epi_tile) {
TiledMma tiled_mma;
auto acc_shape = partition_accumulator_shape();
// ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,ACC_PIPE) where ACC_PIPE=2 so we can double buffer our accumulators for mainloop and epilogue.
Tensor accumulators = cutlass::detail::make_sm100_accumulator<AccumulatorPipelineStageCount, IsOverlappingAccum>(
tiled_mma, acc_shape, EpilogueTile{});
TmemStorage<decltype(accumulators)> tmem_storage;
tmem_storage.accumulators = accumulators;
return tmem_storage;
}
template <class TmemStorage>
CUTLASS_DEVICE static
void
set_tmem_offsets(TmemStorage& tmem_storage, uint32_t tmem_base_addr) {
tmem_storage.accumulators.data() = tmem_base_addr;
}
/// Set up the data needed by this collective for load.
/// Return tuple element contain
/// gA_mkl - The tiled tensor for input A
/// gB_nkl - The tiled tensor for input B
/// tAsA - partitioned smem tensor for A
/// tBsB - partitioned smem tensor for B
template <class ProblemShape_MNKL>
CUTLASS_DEVICE auto
load_init_tma(
ProblemShape_MNKL const& problem_shape_MNKL,
TensorStorage& shared_tensors) const {
using X = Underscore;
// Separate out problem shape for convenience
auto [M,N,K,L] = problem_shape_MNKL;
// TMA
Tensor mA_mkl = observed_tma_load_a_->get_tma_tensor(make_shape(M,K,L));
Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M, BLK_K, m, k, l)
ThrMMA cta_mma = TiledMma{}.get_slice(0);
Tensor tCgA_mkl = cta_mma.partition_A(gA_mkl); // (MMA, MMA_M, MMA_K, m, k, l)
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE)
// Define the CTA-in-cluster Layout and Coord
Layout cta_layout_mnk = make_layout(ClusterShape{});
Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMma::AtomThrID{}));
auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(0);
// Project the cta_layout for tma_a along the n-modes
auto [tAgA_mkl, tAsA] = tma_partition(*observed_tma_load_a_,
get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),
group_modes<0,3>(sA), group_modes<0,3>(tCgA_mkl));
return cute::make_tuple(
shape<3>(gA_mkl), // for scheduler
tAgA_mkl, tAsA // for input tensor values
);
}
template <class ProblemShape_MNKL, class TileScheduler>
CUTLASS_DEVICE auto
load_init_cpasync(
ProblemShape_MNKL const& problem_shape_MNKL,
Params const& params,
TensorStorage& shared_tensors,
TileScheduler const& scheduler,
typename TileScheduler::WorkTileInfo const& work_tile_info) const {
using X = Underscore;
// Separate out problem shape for convenience
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)
// Partition for cpasync
Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
// Build the coordinate tensors with the same shape as input matrices
Tensor cB_nk = make_identity_tensor(make_shape(N,K));
// Slice the coordinate tensors in the same way as A/B tensor partitioning
Tensor cgB_nk = local_tile(cB_nk, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), LoadSmemLayoutB{});
GmemTiledCopyB gmem_to_smem_b_tiled_copy;
int thread_idx = threadIdx.x % NumLoadThreadsCpAsync;
auto thr_copy_b = gmem_to_smem_b_tiled_copy.get_slice(thread_idx);
return cute::make_tuple(
gB_nkl, cgB_nk, sB,
gmem_to_smem_b_tiled_copy, thr_copy_b);
}
/// Set up the data needed by this collective for mma compute.
template <class TmemStorage>
CUTLASS_DEVICE auto
mma_init(
Params const& params,
[[maybe_unused]] TmemStorage tmem_storage,
// [[maybe_unused]] cute::tuple<cute::Tensor<FrgEngine, FrgLayout>, cute::Tensor<FrgEngine, FrgLayout>> const& accumulators_pair,
TensorStorage& shared_tensors) const {
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), MmaSmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
// Allocate "fragments/descriptors" for A and B matrices
Tensor tCrA = TiledMma::make_fragment_A(sA); // (MMA,MMA_M,MMA_K,PIPE)
Tensor tCrB = TiledMma::make_fragment_B(sB); // (MMA,MMA_N,MMA_K,PIPE)
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<3>(sA)); // PIPE
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<3>(sB));
TiledMma tiled_mma;
if constexpr (IsRuntimeDataType) {
// Update instruction descriptor according to runtime argument.
// Applying bitmask (0b111) to help compiler deduce that the conversion and assignment are safe.
tiled_mma.idesc_.a_format_ = uint8_t(params.runtime_data_type_a) & 0b111;
tiled_mma.idesc_.b_format_ = uint8_t(params.runtime_data_type_b) & 0b111;
}
return cute::make_tuple(tiled_mma, tCrA, tCrB);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Producer Perspective
template <
class KTileCount,
class GTensorPartitionedA,
class STensorA,
class TileCoordMNKL,
class KTileIterator
>
CUTLASS_DEVICE auto
load_tma(
MainloopPipelineTMA mainloop_pipeline,
MainloopPipelineTMAState mainloop_pipe_producer_state,
cute::tuple<KTileCount,
GTensorPartitionedA,
STensorA> const& load_inputs,
TileCoordMNKL const& cta_coord_mnkl,
KTileIterator k_tile_iter, int k_tile_count) {
// Unpack from load_inputs
KTileCount k_tiles = get<0>(load_inputs);
GTensorPartitionedA tAgA_mkl = get<1>(load_inputs);
STensorA tAsA = get<2>(load_inputs);
// slice out the work coord from partitioned tensors
Tensor tAgA = tAgA_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl));
auto barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state);
// Issue the Mainloop loads
CUTLASS_PRAGMA_NO_UNROLL
while (k_tile_count > 0) {
// LOCK mainloop_pipe_producer_state for _writing_
mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state, barrier_token);
using BarrierType = typename MainloopPipelineTMA::ProducerBarrierType;
BarrierType* tma_barrier = mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state);
int write_stage = mainloop_pipe_producer_state.index();
++mainloop_pipe_producer_state;
barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state);
if (cute::elect_one_sync()) {
copy(observed_tma_load_a_->with(*tma_barrier), tAgA(_,*k_tile_iter), tAsA(_,write_stage));
}
--k_tile_count;
++k_tile_iter;
}
return cute::make_tuple(mainloop_pipe_producer_state, k_tile_iter);
}
template <
// class GTensorB,
// class CTensorB,
// class STensorB,
// class ProblemShape_MNKL,
// class TiledCopyB,
// class ThreadCopyB,
class TileCoordMNKL,
class KTileIterator,
class ProblemShape_MNKL,
class... TParams
>
CUTLASS_DEVICE auto
load_cpasync(
Params const& params,
MainloopPipelineCpAsync mainloop_pipeline,
MainloopPipelineCpAsyncState mainloop_pipe_producer_state,
cute::tuple<TParams...> const& load_inputs,
TileCoordMNKL const& cta_coord_mnkl,
KTileIterator k_tile_iter, int k_tile_count,
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,
gmem_to_smem_b_tiled_copy, thr_copy_b] = load_inputs;
auto [M,N,K,L] = effective_shape;
// Slice out the work coord from partitioned tensors
Tensor gB_in = tBgB_nkl(_, _, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
// Repeat slicing out coordinate tensor exactly the same as input tensor does
Tensor cgB_nk_in = cgB_nk(_, _, get<1>(cta_coord_mnkl), _);
auto k_residue = K - size<1>(gB_in) * size<2>(gB_in); // K - BLK_K * k is negative
Tensor gB = gB_in;
Tensor cB = cgB_nk_in;
auto tBgB = thr_copy_b.partition_S(gB);
auto tBsB = thr_copy_b.partition_D(sB);
// Allocate predicate tensors for n
Tensor tBpB = make_tensor<bool>(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{});
Tensor tBcB_nk = thr_copy_b.partition_S(cgB_nk_in);
Tensor tBcB = thr_copy_b.partition_S(cB);
// Copy gmem to smem for *k_tile_iter, predicating for k residue
Tensor tBgBk = tBgB(_,_,_,*k_tile_iter);
// Repeating on predicators with the same operations on tBgB
Tensor tBcBk = tBcB(_,_,_,*k_tile_iter);
// Set predicates for n bounds
CUTLASS_PRAGMA_UNROLL
for (int n = 0; n < size<0>(tBpB); ++n) {
tBpB(n,0) = elem_less(get<0>(tBcBk(0,n,0)), N); // blk_n coord < N
}
// we will process the last tile after the mainloop
if (k_residue != 0) {
--k_tile_count;
}
// Issue the Mainloop loads
CUTLASS_PRAGMA_NO_UNROLL
while (k_tile_count > 0) {
mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state);
int write_stage = mainloop_pipe_producer_state.index();
copy_if(gmem_to_smem_b_tiled_copy, tBpB, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
mainloop_pipeline.producer_commit(mainloop_pipe_producer_state, cutlass::arch::cpasync_barrier_arrive);
--k_tile_count;
++k_tile_iter;
++mainloop_pipe_producer_state;
}
// last tile with predication on k to account for residue
// For performance consideration,
// this predicated block for K-tail is only activated when there is k-residue
if (k_residue != 0) {
// LOCK mainloop_pipe_producer_state for _writing_
mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state);
int write_stage = mainloop_pipe_producer_state.index();
CUTLASS_PRAGMA_UNROLL
for (int k = 0; k < size<2>(tBsB); ++k) {
if (int(get<1>(tBcBk(0,0,k))) >= 0) { // blk_k coord < K
copy_if(gmem_to_smem_b_tiled_copy, tBpB(_,k), tBgB(_,_,k,*k_tile_iter), tBsB(_,_,k,write_stage));
}
else {
clear(tBsB(_,_,k,write_stage));
}
}
++k_tile_iter;
--k_tile_count;
// UNLOCK mainloop_pipe_producer_state
mainloop_pipeline.producer_commit(mainloop_pipe_producer_state, cutlass::arch::cpasync_barrier_arrive);
// Advance mainloop_pipe_producer_state
++mainloop_pipe_producer_state;
}
return cute::make_tuple(mainloop_pipe_producer_state, k_tile_iter);
}
/// Perform a Producer Epilogue to prevent early exit of ctas in a Cluster
CUTLASS_DEVICE void
load_tail_tma(MainloopPipelineTMA mainloop_pipeline, MainloopPipelineTMAState mainloop_pipe_producer_state) {
// Issue the epilogue waits
// This helps avoid early exit of ctas in Cluster
// Waits for all stages to either be released (all
// Consumer UNLOCKs), or if the stage was never used
// then would just be acquired since the phase was
// still inverted from make_producer_start_state
mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);
}
CUTLASS_DEVICE void
load_tail_cpasync(MainloopPipelineCpAsync mainloop_pipeline, MainloopPipelineCpAsyncState mainloop_pipe_producer_state) {
mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class AccumulatorPipeline,
class FrgEngine, class FrgLayout,
class FragmentA, class FragmentB,
class CtaTileCoord
>
CUTLASS_DEVICE auto
mma(cute::tuple<MainloopPipelineTMA,
MainloopPipelineCpAsync,
AccumulatorPipeline> pipelines,
cute::tuple<MainloopPipelineTMAState,
MainloopPipelineCpAsyncState,
typename AccumulatorPipeline::PipelineState> pipeline_states,
cute::tuple<cute::Tensor<FrgEngine, FrgLayout>> const& accumulators_pair,
cute::tuple<TiledMma, FragmentA, FragmentB> const& mma_inputs,
CtaTileCoord cta_tile_coord,
int k_tile_count
) {
static_assert(is_tmem<FrgEngine>::value, "Accumulator must be tmem resident.");
static_assert(rank(FrgLayout{}) == 3, "Accumulator must be MMA-partitioned: (MMA, MMA_M, MMA_N)");
auto accumulators = get<0>(accumulators_pair);
auto [tiled_mma, tCrA, tCrB] = mma_inputs;
auto [mainloop_pipeline_tma, mainloop_pipeline_cpasync, accumulator_pipeline] = pipelines;
auto [mainloop_pipe_tma_consumer_state, mainloop_pipe_cpasync_consumer_state, accumulator_pipe_producer_state] = pipeline_states;
//
// PIPELINED MAIN LOOP
//
tiled_mma.accumulate_ = UMMA::ScaleOut::Zero;
// Wait for tmem accumulator buffer to become empty with a flipped phase
accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);
CUTLASS_PRAGMA_NO_UNROLL
while (k_tile_count > 0) {
mainloop_pipeline_tma.consumer_wait(mainloop_pipe_tma_consumer_state);
mainloop_pipeline_cpasync.consumer_wait(mainloop_pipe_cpasync_consumer_state);
int read_stage_tma = mainloop_pipe_tma_consumer_state.index();
int read_stage_cpasync = mainloop_pipe_cpasync_consumer_state.index();
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_tma), tCrB(_,_,k_block,read_stage_cpasync), accumulators);
tiled_mma.accumulate_ = UMMA::ScaleOut::One;
}
mainloop_pipeline_tma.consumer_release(mainloop_pipe_tma_consumer_state);
mainloop_pipeline_cpasync.consumer_release(mainloop_pipe_cpasync_consumer_state);
--k_tile_count;
++mainloop_pipe_tma_consumer_state;
++mainloop_pipe_cpasync_consumer_state;
}
return cute::make_tuple(mainloop_pipe_tma_consumer_state, mainloop_pipe_cpasync_consumer_state);
}
protected:
typename Params::TMA_A const* observed_tma_load_a_{nullptr};
RuntimeDataTypeA runtime_data_type_a_{};
RuntimeDataTypeB runtime_data_type_b_{};
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/sm100_mma_warpspecialized_mixed_input.hpp
+3 -3
View File
@@ -248,7 +248,7 @@ public:
using Load2MmaPipeline = cutlass::PipelineTmaUmmaAsync<
DispatchPolicy::Load2TransformPipelineStageCount,
ClusterShape,
ClusterShape,
AtomThrShapeMNK>;
using Load2MmaPipelineState = typename Load2MmaPipeline::PipelineState;
@@ -316,7 +316,7 @@ public:
using SmemLayoutACompute = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomACompute{},
append(CtaShapeA_MK{}, Int<DispatchPolicy::Load2TransformPipelineStageCount>{}),
append(CtaShapeA_MK{}, Int<DispatchPolicy::Transform2MmaPipelineStageCount>{}),
(cute::conditional_t<cutlass::gemm::detail::is_mn_major<StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})));
using SmemLayoutB = decltype(UMMA::tile_to_mma_shape(
@@ -385,7 +385,7 @@ public:
struct TensorStorageUntransformed {
alignas(512) cute::ArrayEngine<ElementA, cute::cosize_v<SmemLayoutA>> smem_A;
cute::ArrayEngine<ElementB, cute::cosize_v<SmemLayoutB>> smem_B;
alignas(1024) cute::ArrayEngine<ElementB, cute::cosize_v<SmemLayoutB>> smem_B;
cute::ArrayEngine<NonVoidElementScale, scale_elements> smem_scale;
cute::ArrayEngine<NonVoidElementZero, zero_elements> smem_zero;
};
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+11 -14
View File
@@ -73,7 +73,7 @@ template <
class SmemCopyAtomB_,
class TransformB_>
struct CollectiveMma<
MainloopSm90ArrayTmaGmmaWarpSpecializedBlockScaling<Stages, ClusterShape, KernelSchedule>,
MainloopSm90ArrayTmaGmmaWarpSpecializedBlockwise<Stages, ClusterShape, KernelSchedule>,
TileShape_,
ElementA_,
StridePairA_,
@@ -92,7 +92,7 @@ struct CollectiveMma<
//
// Type Aliases
//
using DispatchPolicy = MainloopSm90ArrayTmaGmmaWarpSpecializedBlockScaling<Stages, ClusterShape, KernelSchedule>;
using DispatchPolicy = MainloopSm90ArrayTmaGmmaWarpSpecializedBlockwise<Stages, ClusterShape, KernelSchedule>;
using TileShape = TileShape_;
using ElementA = ElementA_;
using StrideA = cute::tuple_element_t<0,StridePairA_>;
@@ -382,8 +382,6 @@ struct CollectiveMma<
auto [M,N,K,L] = problem_shape_MNKL;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), InternalStrideA{});
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), InternalStrideB{});
// We expect full tiles in K
implementable = implementable && K % size<2>(TileShape{}) == 0;
}
}
@@ -824,16 +822,13 @@ struct CollectiveMma<
// Prologue GMMAs
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
// 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);
// fence_operand();
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
@@ -977,7 +972,7 @@ struct CollectiveMma<
++smem_pipe_release;
}
if (k_tile_count) {
if (k_tile_count > 0) {
pipeline.consumer_wait(smem_pipe_read, barrier_token);
//
@@ -1072,9 +1067,11 @@ struct CollectiveMma<
/// 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);
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);
}
}
//
include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+2 -8
View File
@@ -73,7 +73,7 @@ template <
class SmemCopyAtomB_,
class TransformB_>
struct CollectiveMma<
MainloopSm90TmaGmmaWarpSpecializedBlockScalingFP8<Stages, ClusterShape, KernelSchedule>,
MainloopSm90TmaGmmaWarpSpecializedBlockwiseFP8<Stages, ClusterShape, KernelSchedule>,
TileShape_,
ElementA_,
StridePairA_,
@@ -91,7 +91,7 @@ struct CollectiveMma<
//
// Type Aliases
//
using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecializedBlockScalingFP8<Stages, ClusterShape, KernelSchedule>;
using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecializedBlockwiseFP8<Stages, ClusterShape, KernelSchedule>;
using TileShape = TileShape_;
using ElementA = ElementA_;
using StrideA = cute::tuple_element_t<0,StridePairA_>;
@@ -391,12 +391,6 @@ struct CollectiveMma<
implementable = false;
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem size doesn't meet the minimum alignment requirements for using TMA to load scale B.\n");
}
// We expect full tiles in K
if (K % size<2>(TileShape{}) != 0) {
implementable = false;
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem size K is incompatible with tile size.\n");
}
return implementable;
}
include/cutlass/gemm/dispatch_policy.hpp
+60 -15
View File
@@ -127,10 +127,15 @@ struct KernelPtrArrayTmaWarpSpecializedCooperative { };
struct KernelPtrArrayTmaWarpSpecializedPingpong { };
// FP8 related policies (including Blocked Scaled Accumulation)
struct KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum: KernelTmaWarpSpecializedCooperative { };
struct KernelTmaWarpSpecializedPingpongFP8BlockScaledAccum: KernelTmaWarpSpecializedPingpong { };
struct KernelPtrArrayTmaWarpSpecializedCooperativeFP8BlockScaledAccum: KernelPtrArrayTmaWarpSpecializedCooperative { };
struct KernelPtrArrayTmaWarpSpecializedPingpongFP8BlockScaledAccum: KernelPtrArrayTmaWarpSpecializedPingpong { };
struct KernelTmaWarpSpecializedCooperativeFP8Blockwise: KernelTmaWarpSpecializedCooperative { };
struct KernelTmaWarpSpecializedPingpongFP8Blockwise: KernelTmaWarpSpecializedPingpong { };
struct KernelPtrArrayTmaWarpSpecializedCooperativeFP8Blockwise: KernelPtrArrayTmaWarpSpecializedCooperative { };
struct KernelPtrArrayTmaWarpSpecializedPingpongFP8Blockwise: KernelPtrArrayTmaWarpSpecializedPingpong { };
using KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum = KernelTmaWarpSpecializedCooperativeFP8Blockwise;
using KernelTmaWarpSpecializedPingpongFP8BlockScaledAccum = KernelTmaWarpSpecializedPingpongFP8Blockwise;
using KernelPtrArrayTmaWarpSpecializedCooperativeFP8BlockScaledAccum = KernelPtrArrayTmaWarpSpecializedCooperativeFP8Blockwise;
using KernelPtrArrayTmaWarpSpecializedPingpongFP8BlockScaledAccum = KernelPtrArrayTmaWarpSpecializedPingpongFP8Blockwise;
// Policies to opt into mixed type GEMMs
struct KernelTmaWarpSpecializedMixedInput : KernelTmaWarpSpecialized { };
@@ -319,17 +324,17 @@ struct MainloopSm90TmaGmmaWarpSpecializedFP8
// n-buffer in smem (Hopper TMA), pipelined with Hopper GMMA and TMA, Warp specialized dynamic schedule
// For FP8 kernels with Block Scaling
// For FP8 kernels with Blockwise (Software) Scaling
template<
int Stages_,
class ClusterShape_ = Shape<_1,_1,_1>,
class KernelSchedule = KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum
class KernelSchedule = KernelTmaWarpSpecializedCooperativeFP8Blockwise
>
struct MainloopSm90TmaGmmaWarpSpecializedBlockScalingFP8
struct MainloopSm90TmaGmmaWarpSpecializedBlockwiseFP8
: MainloopSm90TmaGmmaWarpSpecialized<Stages_, ClusterShape_, KernelSchedule> {
static_assert(
cute::is_same_v<KernelSchedule, KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum> ||
cute::is_same_v<KernelSchedule, KernelTmaWarpSpecializedPingpongFP8BlockScaledAccum>,
cute::is_same_v<KernelSchedule, KernelTmaWarpSpecializedCooperativeFP8Blockwise> ||
cute::is_same_v<KernelSchedule, KernelTmaWarpSpecializedPingpongFP8Blockwise>,
"KernelSchedule must be one of the warp specialized FP8 block scale policies");
};
@@ -411,15 +416,15 @@ struct MainloopSm90ArrayTmaGmmaWarpSpecializedMixedInput {
template<
int Stages_,
class ClusterShape_ = Shape<_1,_1,_1>,
class KernelSchedule = KernelPtrArrayTmaWarpSpecializedCooperativeFP8BlockScaledAccum
class KernelSchedule = KernelPtrArrayTmaWarpSpecializedCooperativeFP8Blockwise
>
struct MainloopSm90ArrayTmaGmmaWarpSpecializedBlockScaling
struct MainloopSm90ArrayTmaGmmaWarpSpecializedBlockwise
: MainloopSm90ArrayTmaGmmaWarpSpecialized<Stages_, ClusterShape_, KernelSchedule> {
static_assert(
cute::is_any_of_v<
KernelSchedule,
KernelPtrArrayTmaWarpSpecializedCooperativeFP8BlockScaledAccum,
KernelPtrArrayTmaWarpSpecializedPingpongFP8BlockScaledAccum
KernelPtrArrayTmaWarpSpecializedCooperativeFP8Blockwise,
KernelPtrArrayTmaWarpSpecializedPingpongFP8Blockwise
>,
"KernelSchedule must be one of the warp specialized FP8 block scale policies");
};
@@ -440,6 +445,15 @@ struct KernelWarpSpecializedSm100 final {
static constexpr int AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;
};
template<
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_
>
struct KernelMixedTmaCpAsyncWarpSpecializedSm100 final {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
static constexpr int AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;
};
template<
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_
@@ -653,7 +667,7 @@ template<
class KernelSchedule
>
struct HasAuxiliaryLoad<
MainloopSm90ArrayTmaGmmaWarpSpecializedBlockScaling<
MainloopSm90ArrayTmaGmmaWarpSpecializedBlockwise<
Stages,
ClusterShape,
KernelSchedule
@@ -666,7 +680,7 @@ template<
class KernelSchedule
>
struct HasAuxiliaryLoad<
MainloopSm90TmaGmmaWarpSpecializedBlockScalingFP8<
MainloopSm90TmaGmmaWarpSpecializedBlockwiseFP8<
Stages,
ClusterShape,
KernelSchedule
@@ -700,6 +714,7 @@ struct KernelScheduleSm100DenseGemm : KernelScheduleSm100 {}; // Base policy
struct KernelTmaWarpSpecialized1SmSm100 final : KernelSchedule1Sm, KernelScheduleSm100DenseGemm {}; // Use for 1SM Dense GEMM Kernels for Collective Mainloop Builder
struct KernelTmaWarpSpecialized2SmSm100 final : KernelSchedule2Sm, KernelScheduleSm100DenseGemm {}; // Use for 2SM Dense GEMM Kernels for Collective Mainloop Builder
struct KernelWarpSpecialized1SmSm100 final : KernelSchedule1Sm, KernelScheduleSm100DenseGemm {}; // Use for 1SM Dense GEMM Kernels for Collective Mainloop Builder Without TMA
struct KernelMixedTmaCpAsyncWarpSpecialized1SmSm100 final : KernelSchedule1Sm, KernelScheduleSm100DenseGemm {};
///////////////////////////////////////////////////////////////////////////////////////////////////////
// SM100 Ptr-Array Dense GEMM Dispatch Policies
@@ -795,6 +810,8 @@ struct KernelTmaWarpSpecialized1SmMxf4Sm100 final : KernelSchedule1
struct KernelTmaWarpSpecialized2SmMxf4Sm100 final : KernelSchedule2Sm, KernelScheduleMxNvf4Sm100 { };
struct KernelTmaWarpSpecialized1SmMxf8f6f4Sm100 final : KernelSchedule1Sm, KernelScheduleMxf8f6f4Sm100 { };
struct KernelTmaWarpSpecialized2SmMxf8f6f4Sm100 final : KernelSchedule2Sm, KernelScheduleMxf8f6f4Sm100 { };
struct KernelMixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100 final : KernelSchedule1Sm, KernelScheduleBlockScaledGemmSm100 {};
///////////////////////////////////////////////////////////////////////////////////////////////////////
// SM100 BlockScaled Ptr Array Dense GEMM Dispatch Policies
///////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -950,6 +967,34 @@ struct MainloopSm100UmmaCpAsyncWarpSpecialized {
using Schedule = KernelWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
>
struct MainloopSm100UmmaMixedTmaCpAsyncWarpSpecialized {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using Schedule = KernelMixedTmaCpAsyncWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
>
struct MainloopSm100UmmaMixedTmaCpAsyncWarpSpecializedBlockScaled {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using Schedule = KernelMixedTmaCpAsyncWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
template<
int Stages_,
include/cutlass/gemm/group_array_problem_shape.hpp
+20
View File
@@ -79,6 +79,16 @@ struct GroupProblemShape {
}
};
template <class ProblemShape_, class MaxProblemShape_>
struct MoEProblemShape {
using UnderlyingProblemShape = ProblemShape_;
using MaxProblemShape = MaxProblemShape_;
UnderlyingProblemShape problem_shape;
MaxProblemShape max_problem_shape;
};
template <class ProblemShape_>
class ArrayProblemShape {
public:
@@ -120,4 +130,14 @@ private:
UnderlyingProblemShape problem_shape_{};
};
namespace detail {
template<class T>
struct is_moe_problem_shape : cute::false_type {};
template<class T, class U>
struct is_moe_problem_shape<cutlass::gemm::MoEProblemShape<T,U>> : cute::true_type {};
}
} // namespace cutlass::gemm
include/cutlass/gemm/kernel/gemm_universal.hpp
+1
View File
@@ -73,6 +73,7 @@ struct IsCutlass3ArrayKernel<ProblemShape, cute::void_t<typename ProblemShape::U
#include "cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_mma_transform.hpp"
#include "cutlass/gemm/kernel/sm100_sparse_gemm_tma_warpspecialized.hpp"
#include "cutlass/gemm/kernel/sm100_gemm_cpasync_warpspecialized.hpp"
#include "cutlass/gemm/kernel/sm100_gemm_mixed_tma_cpasync_warpspecialized.hpp"
#include "cutlass/gemm/kernel/sm103_blockscaled_gemm_tma_warpspecialized.hpp"
#include "cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp"
#include "cutlass/gemm/kernel/sm120_gemm_tma_warpspecialized_cooperative_asymmetric_dma.hpp"
include/cutlass/gemm/kernel/sm100_gemm_mixed_tma_cpasync_warpspecialized.hpp
+1011
View File
File diff suppressed because it is too large Load Diff
include/cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp
+21
View File
@@ -240,6 +240,27 @@ public:
void
fixup(WorkTileInfo const&, FrgTensorC&, uint32_t, uint32_t, uint32_t = 1) const { }
template <
bool IsComplex,
class TiledMma,
class AccEngine,
class AccLayout,
class AccumulatorPipeline,
class AccumulatorPipelineState,
class CopyOpT2R
>
CUTLASS_DEVICE
AccumulatorPipelineState
fixup(
TiledMma const& ,
WorkTileInfo const&,
cute::Tensor<AccEngine, AccLayout>&,
AccumulatorPipeline,
AccumulatorPipelineState acc_pipe_consumer_state,
CopyOpT2R) const {
return acc_pipe_consumer_state;
}
template <class ProblemShape, class ElementAccumulator>
static size_t
get_workspace_size(Arguments const& args, ProblemShape problem_shape, KernelHardwareInfo const& hw_info, uint32_t, uint32_t = 1, uint32_t = 1) {
include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp
+1 -1
View File
@@ -991,7 +991,7 @@ public:
mainloop_sf_pipeline,
mainloop_sf_pipe_producer_state,
load_inputs,
cta_coord_mnkl,
cta_coord_mnk,
k_tile_iter_next, k_tile_count - k_tile_prologue,
false, /* did_batch_change - prologue loads handle tensormap acquire */
enable_prefetch ? k_tile_count - k_tile_prologue : 0
include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp
+1 -4
View File
@@ -831,8 +831,6 @@ public:
collective_epilogue.template tensormaps_fence_acquire<IsEpiLoad>(epi_load_tensormap);
}
bool wait = work_tile_info.is_valid() && curr_batch != next_work_tile_info.L_idx;
epi_load_pipe_producer_state = collective_epilogue.load(
epi_load_pipeline,
epi_load_pipe_producer_state,
@@ -843,8 +841,7 @@ public:
lane_idx,
shared_storage.tensors.epilogue,
epi_load_tensormap,
work_tile_info.reduction_subtile_idx(),
wait
work_tile_info.reduction_subtile_idx()
);
}
include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp
+1 -4
View File
@@ -869,8 +869,6 @@ public:
collective_epilogue.template tensormaps_fence_acquire<IsEpiLoad>(epi_load_tensormap);
}
bool wait = work_tile_info.is_valid() && curr_batch != next_work_tile_info.L_idx;
epi_load_pipe_producer_state = collective_epilogue.load(
epi_load_pipeline,
epi_load_pipe_producer_state,
@@ -881,8 +879,7 @@ public:
lane_idx,
shared_storage.tensors.epilogue,
epi_load_tensormap,
work_tile_info.reduction_subtile_idx(),
wait
work_tile_info.reduction_subtile_idx()
);
}
media/docs/cpp/profiler.md
+54 -1
View File
@@ -79,7 +79,7 @@ Instruction shape levels control the selection of WGMMA shapes used in kernel ge
- **Level 2**: Includes shapes that are powers of 2.
- **Level 3**: Includes all other shapes.
The detailed defination of the three instantiation levels controlling cluster shape, MMA shape multiplier, and instruction shape can be found in [sm90_shapes.py](https://github.com/NVIDIA/cutlass/tree/main/python/cutlass_library/sm90_shapes.py).
The detailed definition of the three instantiation levels controlling cluster shape, MMA shape multiplier, and instruction shape can be found in [sm90_shapes.py](https://github.com/NVIDIA/cutlass/tree/main/python/cutlass_library/sm90_shapes.py).
Schedule pruning levels decide the epilogue schedule and mainloop schedule to stamp out a kernel instance. As defined in `get_valid_schedules` in [sm90_utils.py](https://github.com/NVIDIA/cutlass/tree/main/python/cutlass_library/sm90_utils.py),
@@ -122,6 +122,55 @@ For each mixed dtype kernel, the kernel generator will generate combinations of
For {4-bits-dtype, 8-bits-dtype} x 16-bits-dtype, the kernel generator will further generate kernels using shuffled layouts for the narrow data type matrix, which may have a better performance compared to its non-shuffle counter parts.
## Instantiating more kernels with Blackwell
Blackwell (SM100) and Blackwell Ultra similarly support
`CUTLASS_LIBRARY_INSTANTIATION_LEVEL`, in order to instantiate all possible combinations.
Due to this, `CUTLASS_LIBRARY_KERNELS` must be non-empty, since generating and filtering these
kernels alone can take hours.
You must also exercise caution, because not all of these configs are tested, and some may fail to
compile or fail to launch at runtime.
```bash
$ cmake .. \
-DCUTLASS_NVCC_ARCHS="100f" \
-DCUTLASS_LIBRARY_KERNELS="cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f32_*" \
-DCUTLASS_LIBRARY_INSTANTIATION_LEVEL="max" \
-DCUTLASS_UNITY_BUILD_ENABLED=ON
```
The CUTLASS profiler uses the same four-digit integer level (global instantiation level) mechanism to manage the generation of kernel configurations for Blackwell as well:
0. **Instruction Shape**
1. **MMA Shape Multiplier**
2. **Cluster Shape**
3. **Data Type and Schedule Pruning**
Note for Blackwell kernels an MMA shape multiplier is no longer necessary since Blackwell kernels do not have a different
ping pong or cooperative schedule. The profiler ignores this digit when instantiating.
Cluster shape levels define the number of CTAs (Cooperative Thread Arrays) included in the kernel generation:
- **Level 0**: Only dynamic cluster shapes.
- **Level 1**: For 1SM kernels `(1, 1, 1)` and `(2, 1, 1)` for 2SM kernels.
- **Level 2**: For 1SM kernels we also have `(1, 2, 1)` and for 2SM we have `(2, 2, 1)` and `(4, 1, 1)`.
- **Level 3**: For 1SM kernels we have `(1, 4, 1)` and for 2SM we have `(2, 4, 1)` and `(4, 2, 1)`.
- **Level 4**: For 1SM kernels we have `(4, 4, 1)` and for 2SM we have `(4, 4, 1)`.
- **Level 5**: For 1SM kernels we have `(2, 1, 1)`.
- **Level 6**: For 1SM kernels we have `(2, 2, 1)` and `(4, 1, 1)` and for 2SM kernels we have `(8, 1, 1)`.
- **Level 7**: For 1SM kernels we have `(2, 4, 1)` and `(4, 2, 1)`
- **Level 8**: For 1SM kernels we have `(1, 8, 1)` and `(8, 1, 1)`
Instruction shape levels control the selection of MMA shapes used in kernel generation:
- **Level 0**: Generates the "default" shape only.
- **Level 1**: Includes additional shapes for FP8, FP6, and FP4 as well as MX and NVFP4.
- **Level 2**: Includes small tile shapes.
- **Level 3**: Includes some non-power of 2 shapes.
- **Level 4**: Includes further small tile shapes and non-power of 2 shapes.
- **Level 5**: Includes all shapes.
The detailed definition of the three instantiation levels controlling cluster shape and instruction shape can be found in [sm100_shapes.py](https://github.com/NVIDIA/cutlass/tree/main/python/cutlass_library/sm100_shapes.py).
## CUTLASS Profiler usage
The CUTLASS Profiler usage statement may be obtained by executing `cutlass_profiler --help` and appears as follows.
@@ -577,6 +626,10 @@ cutlass3x_sm90_tensorop_gemm_f16_f16_f16_void_f16_128x128x64_1x1x1_0_nnn_align8_
* `f16_f16_f16_void_f16`: In this case, C type is set to `void`, indicating that residual matrix support
is disabled.
## Further Documentation
For documentation on profiling blockwise and groupwise (software scaled) GEMMs see the [example 81 README](https://github.com/NVIDIA/cutlass/blob/main/examples/81_blackwell_gemm_blockwise/README.md).
# Convolution
The CUTLASS Profiler is capable of executing 2-D and 3-D convolution problems for forwards and backwards
media/docs/pythonDSL/cute_dsl_api.rst
+1
View File
@@ -6,6 +6,7 @@ CuTe DSL API
.. toctree::
:maxdepth: 1
changelog <cute_dsl_api/changelog.rst>
cute <cute_dsl_api/cute.rst>
cute_arch <cute_dsl_api/cute_arch.rst>
cute_nvgpu <cute_dsl_api/cute_nvgpu.rst>
media/docs/pythonDSL/cute_dsl_api/changelog.rst
+54
View File
@@ -0,0 +1,54 @@
======================================
Changelog for CuTe DSL API changes
======================================
`4.2.0 <https://github.com/NVIDIA/cutlass/releases/tag/v4.2.0>`_ (2025-09-15)
==============================================================================
* 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.
* Introduced per op fast-math control for math ops(e.g. ``exp``, ``exp2``, ``log2``, ``log``)
* Introduced ``CopyReduceBulkTensorTileS2GOp`` in `tcgen05/copy.py <https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py>`_ to support TMA Reduce.
`4.1.0 <https://github.com/NVIDIA/cutlass/releases/tag/v4.1.0>`_ (2025-07-16)
==============================================================================
* for loop
- Python built-in ``range`` now always generates codes and executes at runtime
- ``cutlass.range`` is advanced ``range`` with kernel code level unrolling and pipelining control
- Deprecated ``cutlass.range_dynamic``, please replace with ``range`` or ``cutlass.range``
- **Experimental** Added ``pipelining`` control for compiler generated software pipeline code
* while/if
- ``while``/``if`` now by default generates codes and executes at runtime unless ``cutlass.const_expr`` is specified for the predicate
- Deprecated ``cutlass.dynamic_expr``, please remove it
* Rename mbarrier functions to reduce ambiguity
* Modify SyncObject API (``MbarrierArray``, ``NamedBarrier``, ``TmaStoreFence``) to match ``std::barrier``
* Change pipeline ``create`` function to take only keyword arguments, and make ``barrier_storage`` optional.
* Introduce ``cutlass.cute.arch.get_dyn_smem_size`` api to get runtime dynamic shared memory size.
* Various API Support for SM100 BlockScaled Gemm
- Introduce BlockScaled MmaOps in `tcgen05/mma.py <https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py>`_, and provide a ``make_blockscaled_trivial_tiled_mma`` function in `blackwell_helpers.py <https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/utils/blackwell_helpers.py>`_ to help construct a BlockScaled TiledMma.
- Introduce S2T CopyOps in `tcgen05/copy.py <https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py>`_.
- Introduce BlockScaled layout utilities in `blockscaled_layout.py <https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/utils/blockscaled_layout.py>`_ for creating the required scale factor layouts in global memory, shared memory and tensor memory.
* ``cutlass.cute.compile`` now supports compilation options. Refer to `JIT compilation options <https://docs.nvidia.com/cutlass/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.html>`_ for more details.
* ``cutlass.cute.testing.assert_`` now works for device JIT function. Specify ``--enable-device-assertions`` as compilation option to enable.
* ``cutlass.cute.make_tiled_copy`` is now deprecated. Please use ``cutlass.cute.make_tiled_copy_tv`` instead.
* Shared memory capacity query
- Introduce ``cutlass.utils.get_smem_capacity_in_bytes`` for querying the shared memory capacity.
- ``<arch>_utils.SMEM_CAPACITY["<arch_str>"]`` is now deprecated.
`4.0.0 <https://github.com/NVIDIA/cutlass/releases/tag/v4.0.0>`_ (2025-06-03)
==============================================================================
* Fixed API mismatch in class ``cute.runtime.Pointer``: change ``element_type`` to ``dtype`` to match ``typing.Pointer``
media/docs/pythonDSL/cute_dsl_general/dsl_control_flow.rst
+49
View File
@@ -72,6 +72,55 @@ All loop indices must be |Constexpr|.
for i in cutlass.range(bound, unroll=2):
cute.printf("%d\\n", i)
Software Pipelining
~~~~~~~~~~~~~~~~~~~
Software pipelining is a technique used to optimize loops. Typically, this involves writing a prefetch loop and a main loop.
.. code-block:: python
@cute.jit
def example():
...
# build a circular buffer
buffer = ...
# prefetch loop
for i in range(prefetch_stages):
cute.copy(atom, gmem[i], buffer[i], ...)
# main loop
for i in range(bound):
if i + prefetch_stages < bound:
cute.copy(atom, gmem[i + prefetch_stages], buffer[(i + prefetch_stages) % total_stages], ...)
use(buffer[i % total_stages])
...
This can be tedious to write and tune. |DSL| provides a loop attribute to ask the compiler to do this.
.. code-block:: python
@cute.jit
def example():
...
# build a circular buffer
buffer = ...
for i in cutlass.range(bound, prefetch_stages=prefetch_stages):
# Compiler automatically handles the pipelining:
# - Generates prefetch loop for initial stages
# - In main loop, prefetches future data while using current data
cute.copy(atom, gmem[i], buffer[i % total_stages], ...)
use(buffer[i % total_stages]) # Uses data from previous iterations
...
Compiler will automatically generate the prefetch loop with `prefetch_stages` iterations and a corresponding main loop.
This feature is experimental and only supported on sm90 and above.
If-Else Statements
------------------
media/docs/pythonDSL/cute_dsl_general/framework_integration.rst
+83 -1
View File
@@ -7,7 +7,8 @@ Integration with Frameworks
In order to facilitate the integration of CUTLASS Python with popular frameworks, we leverage the
`DLPack protocol <https://github.com/dmlc/dlpack>`_ and transform tensors originating from these
frameworks to CuTe tensors. The present page documents the conventions, the API available to the
user, and provide example code snippets for common usage patterns.
user, and provide example code snippets for common usage patterns. We also provide a section on how to
bypass the DLPack protocol and directly call the JIT function.
Implicit Conversion
-------------------
@@ -396,3 +397,84 @@ The following example demonstrates how to use ``mark_compact_shape_dynamic`` to
mode=0, divisibility=1, stride_order=(2, 1, 3, 0, 4)
)
# The stride_order is not consistent with the layout
Bypass the DLPack Protocol
--------------------------
In certain scenarios, users may wish to bypass the DLPack protocol and invoke the JIT function directly.
This can be accomplished by creating a lightweight JIT wrapper around the existing JIT function,
utilizing ``cute.ptr`` and ``cute.make_tensor`` to pass pointers and construct tensors directly.
Typical use cases for bypassing DLPack include:
1. Users want to call the JIT function directly to avoid the overhead introduced by the DLPack protocol.
2. DLPack canonicalizes the stride of shape-1 dimensions to 1, which may result in incorrect alignment
propagation and affect memory access or performance.
3. DLPack may lack support for some narrow data types.
The following example illustrates how to bypass the DLPack protocol when invoking a JIT function.
Assume we have a pre-defined ``TensorOpGemm`` kernel whose JIT interface expects three
arguments of type ``cute.Tensor``. To enable direct invocation without DLPack, we first define a JIT wrapper
function that accepts ``cute.Pointer`` types as parameters. Within this wrapper, we use ``cute.make_tensor``
to construct tensors from the provided pointers, and then call the ``TensorOpGemm`` kernel as usual.
.. code-block:: python
@cute.jit
def tensor_op_gemm_wrapper(
a_ptr: cute.Pointer,
b_ptr: cute.Pointer,
c_ptr: cute.Pointer,
m: cutlass.Int32,
n: cutlass.Int32,
k: cutlass.Int32,
l: cutlass.Int32,
):
# Assume alignment of shape to call tensorop_gemm example
m = cute.assume(m, divby=8)
n = cute.assume(n, divby=8)
# Torch is row major
a_layout = cute.make_ordered_layout((m, k, l), order=(0, 1, 2))
b_layout = cute.make_ordered_layout((n, k, l), order=(0, 1, 2))
c_layout = cute.make_ordered_layout((m, n, l), order=(1, 0, 2))
mA = cute.make_tensor(a_ptr, layout=a_layout)
mB = cute.make_tensor(b_ptr, layout=b_layout)
mC = cute.make_tensor(c_ptr, layout=c_layout)
# TensorOpGemm is a pre-defined kernel from our example
tensor_op_gemm = TensorOpGemm(
a_ptr.value_type, c_ptr.value_type, cutlass.Float32, (2, 2, 1)
)
tensor_op_gemm(mA, mB, mC)
To pass a PyTorch tensor to this new JIT wrapper, we retrieve the raw pointer from the PyTorch tensor
and create a ``cute.Pointer`` instance using ``cute.make_ptr``.
This approach allows us to bypass the DLPack protocol entirely, avoiding its overhead and potential
issues with shape-1 dimension handling.
.. code-block:: python
a = torch.randn(
m, k, l, dtype=torch.float16, device="cuda"
).permute(2, 1, 0)
b = torch.randn(
n, k, l, dtype=torch.float16, device="cuda"
).permute(2, 1, 0)
c = torch.randn(
n, m, l, dtype=torch.float16, device="cuda"
).permute(1, 2, 0)
# from cutlass.cute.runtime import make_ptr
a_ptr = make_ptr(
cutlass.Float16, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
b_ptr = make_ptr(
cutlass.Float16, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
c_ptr = make_ptr(
cutlass.Float16, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
tensor_op_gemm_wrapper(a_ptr, b_ptr, c_ptr, m, n, k, l)
pyproject.toml
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "nvidia-cutlass"
version = "4.0.0.0"
version = "4.2.0.0"
description = "CUTLASS"
readme = "README.md"
requires-python = ">=3.8"
python/CuTeDSL/base_dsl/_mlir_helpers/arith.py
+12 -12
View File
@@ -207,7 +207,7 @@ def int_to_int(a, dst_elem_type, *, loc=None, ip=None):
if dst_width == src_width:
return a
elif src_signed and not dst_signed:
elif src_signed != False and not dst_signed:
# Signed -> Unsigned
if dst_width > src_width:
return arith.extui(dst_mlir_type, a, loc=loc, ip=ip)
@@ -216,7 +216,7 @@ def int_to_int(a, dst_elem_type, *, loc=None, ip=None):
elif src_signed == dst_signed:
# Same signedness
if dst_width > src_width:
if src_signed and src_width > 1:
if src_signed != False and src_width > 1:
return arith.extsi(dst_mlir_type, a, loc=loc, ip=ip)
else:
return arith.extui(dst_mlir_type, a, loc=loc, ip=ip)
@@ -479,7 +479,7 @@ class ArithValue(ir.Value):
if self.is_float:
q = arith.divf(self, other, loc=loc, ip=ip)
return math.floor(q, loc=loc, ip=ip)
elif self.signed:
elif self.signed != False:
return arith.floordivsi(self, other, loc=loc, ip=ip)
else:
return arith.divui(self, other, loc=loc, ip=ip)
@@ -489,7 +489,7 @@ class ArithValue(ir.Value):
def __mod__(self, other, *, loc=None, ip=None) -> "ArithValue":
if self.is_float:
return arith.remf(self, other, loc=loc, ip=ip)
elif self.signed:
elif self.signed != False:
return arith.remsi(self, other, loc=loc, ip=ip)
else:
return arith.remui(self, other, loc=loc, ip=ip)
@@ -524,7 +524,7 @@ class ArithValue(ir.Value):
def __lt__(self, other, *, loc=None, ip=None) -> "ArithValue":
if self.is_float:
return arith.cmpf(arith.CmpFPredicate.OLT, self, other, loc=loc, ip=ip)
elif self.signed:
elif self.signed != False:
return arith.cmpi(arith.CmpIPredicate.slt, self, other, loc=loc, ip=ip)
else:
return arith.cmpi(arith.CmpIPredicate.ult, self, other, loc=loc, ip=ip)
@@ -534,7 +534,7 @@ class ArithValue(ir.Value):
def __le__(self, other, *, loc=None, ip=None) -> "ArithValue":
if self.is_float:
return arith.cmpf(arith.CmpFPredicate.OLE, self, other, loc=loc, ip=ip)
elif self.signed:
elif self.signed != False:
return arith.cmpi(arith.CmpIPredicate.sle, self, other, loc=loc, ip=ip)
else:
return arith.cmpi(arith.CmpIPredicate.ule, self, other, loc=loc, ip=ip)
@@ -561,7 +561,7 @@ class ArithValue(ir.Value):
def __gt__(self, other, *, loc=None, ip=None) -> "ArithValue":
if self.is_float:
return arith.cmpf(arith.CmpFPredicate.OGT, self, other, loc=loc, ip=ip)
elif self.signed:
elif self.signed != False:
return arith.cmpi(arith.CmpIPredicate.sgt, self, other, loc=loc, ip=ip)
else:
return arith.cmpi(arith.CmpIPredicate.ugt, self, other, loc=loc, ip=ip)
@@ -571,7 +571,7 @@ class ArithValue(ir.Value):
def __ge__(self, other, *, loc=None, ip=None) -> "ArithValue":
if self.is_float:
return arith.cmpf(arith.CmpFPredicate.OGE, self, other, loc=loc, ip=ip)
elif self.signed:
elif self.signed != False:
return arith.cmpi(arith.CmpIPredicate.sge, self, other, loc=loc, ip=ip)
else:
return arith.cmpi(arith.CmpIPredicate.uge, self, other, loc=loc, ip=ip)
@@ -599,7 +599,7 @@ class ArithValue(ir.Value):
@_dispatch_to_rhs_r_op
@_binary_op
def __rshift__(self, other, *, loc=None, ip=None) -> "ArithValue":
if self.signed:
if self.signed != False:
return arith.shrsi(self, other, loc=loc, ip=ip)
else:
return arith.shrui(self, other, loc=loc, ip=ip)
@@ -633,7 +633,7 @@ class ArithValue(ir.Value):
return super().__hash__()
def __str__(self):
return super().__str__().replace(ir.Value.__name__, ArithValue.__name__)
return "?"
def __repr__(self):
return self.__str__()
@@ -657,7 +657,7 @@ def _min(lhs, rhs, *, loc=None, ip=None):
rhs = arith.constant(lhs.type, rhs, loc=loc, ip=ip)
if arith._is_integer_like_type(lhs.type):
if lhs.signed:
if lhs.signed != False:
return arith.minsi(lhs, rhs, loc=loc, ip=ip)
else:
return arith.minui(lhs, rhs, loc=loc, ip=ip)
@@ -683,7 +683,7 @@ def _max(lhs, rhs, *, loc=None, ip=None):
rhs = arith.constant(lhs.type, rhs, loc=loc, ip=ip)
if arith._is_integer_like_type(lhs.type):
if lhs.signed:
if lhs.signed != False:
return arith.maxsi(lhs, rhs, loc=loc, ip=ip)
else:
return arith.maxui(lhs, rhs, loc=loc, ip=ip)
python/CuTeDSL/base_dsl/ast_helpers.py
+132 -76
View File
@@ -17,12 +17,16 @@ The preprocessor read through python's ast and changes the input code.
from typing import Callable, Iterator, Optional, overload
from typing_extensions import deprecated
import warnings
import inspect
from types import BuiltinFunctionType
from functools import lru_cache
from .utils.logger import log
from .common import *
from ._mlir_helpers.arith import ArithValue
class Executor:
"""
The Executor class handles dynamic and compile-time (constexpr) execution
@@ -45,9 +49,11 @@ class Executor:
self._compare_executor = None
self._any_executor = None
self._all_executor = None
self._builtin_redirector = None
def set_functions(
self,
*,
is_dynamic_expression: Callable,
loop_execute_range_dynamic: Callable,
if_dynamic: Callable,
@@ -55,6 +61,7 @@ class Executor:
compare_executor: Callable,
any_executor: Callable = None,
all_executor: Callable = None,
builtin_redirector: Callable = None,
):
self._is_dynamic_expression = is_dynamic_expression
self._loop_execute_range_dynamic = loop_execute_range_dynamic
@@ -63,6 +70,7 @@ class Executor:
self._compare_executor = compare_executor
self._any_executor = any_executor
self._all_executor = all_executor
self._builtin_redirector = builtin_redirector
@staticmethod
def convert_to_list(x):
@@ -90,42 +98,18 @@ class Executor:
return res[0]
return res
@staticmethod
def for_constexpr(
func: Callable,
start: int,
stop: int,
step: int,
used_args: list,
iter_args: list,
):
log().debug("start [%s] stop [%s] step [%s]", start, stop, step)
loop_results = iter_args
log().debug("iter_args [%s]", iter_args)
for i in range(start, stop, step):
log().debug("i [%s] iter_args [%s]", i, iter_args)
loop_results = func(i, *used_args, *loop_results)
log().debug("loop_results [%s]", loop_results)
if loop_results is None:
loop_results = []
if not isinstance(loop_results, list):
loop_results = [loop_results]
log().debug("done loop_results [%s]", loop_results)
return Executor.converge_ret_val(loop_results)
def for_execute(
self,
func,
start,
stop,
step,
used_args=[],
iter_args=[],
iter_arg_names=[],
write_args=[],
full_write_args_count=0,
write_args_names=[],
unroll=-1,
unroll_full=False,
pipelining=None,
prefetch_stages=None,
):
assert (
self._loop_execute_range_dynamic
@@ -137,12 +121,12 @@ class Executor:
start,
stop,
step,
used_args,
iter_args,
iter_arg_names,
write_args,
full_write_args_count,
write_args_names,
unroll,
unroll_full,
pipelining,
prefetch_stages,
)
def if_execute(
@@ -150,15 +134,20 @@ class Executor:
pred,
then_block: Callable,
else_block: Optional[Callable] = None,
used_args=[],
yield_args=[],
yield_arg_names=[],
write_args=[],
full_write_args_count=0,
write_args_names=[],
):
assert self._if_dynamic, "Functions must be set before execution."
# MLIR generation
return self._if_dynamic(
pred, then_block, else_block, used_args, yield_args, yield_arg_names
pred,
then_block,
else_block,
write_args,
full_write_args_count,
write_args_names,
)
def while_execute(
@@ -166,9 +155,9 @@ class Executor:
pred,
while_before_block: Callable,
while_after_block: Callable,
used_args=[],
yield_args=[],
yield_arg_names=[],
write_args=[],
full_write_args_count=0,
write_args_names=[],
):
assert self._while_dynamic, "Functions must be set before execution."
@@ -176,9 +165,9 @@ class Executor:
return self._while_dynamic(
while_before_block,
while_after_block,
used_args,
yield_args,
yield_arg_names,
write_args,
full_write_args_count,
write_args_names,
)
@@ -194,23 +183,24 @@ def loop_selector(
stop,
step,
*,
used_args=[],
iter_args=[],
iter_arg_names=[],
write_args=[],
full_write_args_count=0,
write_args_names=[],
unroll=-1,
unroll_full=False,
pipelining=None,
prefetch_stages=None,
):
log().debug(
"start [%s] stop [%s] step [%s] used_args [%s] iter_args [%s] unroll [%s] unroll_full [%s] pipelining [%s]",
"start [%s] stop [%s] step [%s] write_args [%s] full_write_args_count [%s] write_args_names [%s] unroll [%s] unroll_full [%s] prefetch_stages [%s]",
start,
stop,
step,
used_args,
iter_args,
write_args,
full_write_args_count,
write_args_names,
unroll,
unroll_full,
pipelining,
prefetch_stages,
)
from .typing import Integer, Numeric
@@ -230,19 +220,19 @@ def loop_selector(
start,
stop,
step,
used_args,
iter_args,
iter_arg_names,
write_args,
full_write_args_count,
write_args_names,
unroll,
unroll_full,
pipelining,
prefetch_stages,
)
return ir_loop
def if_selector(pred, used_args=[], yield_args=[]):
log().debug("pred [%s] used_args [%s] yield_args [%s]", pred, used_args, yield_args)
def if_selector(pred, write_args=[]):
log().debug("pred [%s] write_args [%s]", pred, write_args)
# Handle Numeric types here?
from .typing import Numeric
@@ -251,14 +241,14 @@ def if_selector(pred, used_args=[], yield_args=[]):
pred = pred.value
def ir_loop(func):
return func(pred, *used_args, *yield_args)
return func(pred, *write_args)
return ir_loop
def while_selector(pred, used_args=[], yield_args=[]):
def while_selector(pred, write_args=[]):
def ir_while_loop(func):
return func(pred, *used_args, *yield_args)
return func(pred, *write_args)
return ir_while_loop
@@ -267,17 +257,17 @@ def while_executor(
pred,
while_before_block: Callable,
while_after_block: Callable,
used_args=[],
yield_args=[],
yield_arg_names=[],
write_args=[],
full_write_args_count=0,
write_args_names=[],
):
return executor.while_execute(
pred,
while_before_block,
while_after_block,
used_args,
yield_args,
yield_arg_names,
write_args,
full_write_args_count,
write_args_names,
)
@@ -285,12 +275,17 @@ def if_executor(
pred,
then_block: Callable,
else_block: Optional[Callable] = None,
used_args=[],
yield_args=[],
yield_arg_names=[],
write_args=[],
full_write_args_count=0,
write_args_names=[],
):
return executor.if_execute(
pred, then_block, else_block, used_args, yield_args, yield_arg_names
pred,
then_block,
else_block,
write_args,
full_write_args_count,
write_args_names,
)
@@ -313,14 +308,17 @@ class range:
- unroll: Number of iterations to unroll (0 or 1 = no unrolling)
- unroll_full: Whether to fully unroll the loop
- pipelining: Compiler generated pipeline configuration
- prefetch_stages: Number of prefetch stages to generate
"""
@overload
def __new__(cls, stop, unroll=0, unroll_full=False, pipelining=None):
def __new__(cls, stop, unroll=0, unroll_full=False, prefetch_stages=None):
pass
@overload
def __new__(cls, start, stop, step, unroll=0, unroll_full=False, pipelining=None):
def __new__(
cls, start, stop, step, unroll=0, unroll_full=False, prefetch_stages=None
):
pass
def __new__(cls, *args, **kwargs):
@@ -340,6 +338,7 @@ def range_dynamic(*args, **kwargs):
def range_constexpr(*args):
raise DSLRuntimeError("range_constexpr should be preprocessed by preprocessor.")
# =============================================================================
# If expressions
# =============================================================================
@@ -405,7 +404,7 @@ def assert_executor(test, msg=None):
else:
raise DSLRuntimeError(
"Only constexpr (Python Value) is allowed here, but got non-constexpr (IR Values) expression.",
suggestion = "Please replace with runtime assert."
suggestion="Please replace with runtime assert.",
)
@@ -413,10 +412,11 @@ def bool_cast(value):
if executor._is_dynamic_expression(value):
raise DSLRuntimeError(
"Only constexpr (Python Value) is allowed here, but got non-constexpr (IR Values) expression.",
suggestion = "Please explicitly convert to boolean with expressions like comparision."
suggestion="Please explicitly convert to boolean with expressions like comparision.",
)
return bool(value)
def compare_executor(left, comparators, ops):
"""
Executes comparison operations with a left operand and a list of comparators.
@@ -470,6 +470,19 @@ def all_executor(iterable):
# =============================================================================
# Control flow checks
# =============================================================================
class DSLOptimizationWarning(Warning):
"""
This warning is used to warn the user about the optimization related issues in DSL.
"""
def __init__(self, message):
self.message = message
super().__init__()
def __str__(self):
return self.message
def range_value_check(*args):
"""
Ensure all `range_constexpr` bounds are compile-time constants (Python ints).
@@ -495,7 +508,7 @@ def range_value_check(*args):
if range_length >= 64:
warnings.warn(
f"This static loop has {range_length} iterations, which may be very slow to compile, consider using `cutlass.range(..., unroll_full=True)` instead.",
category=UserWarning,
category=DSLOptimizationWarning,
stacklevel=2,
)
@@ -519,7 +532,50 @@ def range_perf_warning(filename, lineno, *args):
"This loop is no longer unrolled and may cause performance regression. "
"Use `range(..., unroll_full=True)` for full unrolling, or switch to `range_constexpr` when bounds are compile-time constants."
),
category=UserWarning,
category=DSLOptimizationWarning,
filename=filename,
lineno=lineno,
)
@lru_cache(maxsize=1)
def _get_self_module():
"""
This function is used to get the owning module of this function.
"""
return inspect.getmodule(_get_self_module)
def cf_symbol_check(symbol):
"""
Check if the symbol is control flow symbol from current module.
"""
failed = False
name = symbol.__name__
self_module = _get_self_module()
if inspect.ismodule(symbol):
name = "range"
if not self_module.__name__.startswith(symbol.__name__):
failed = True
else:
owning_module = inspect.getmodule(symbol)
if owning_module != self_module:
failed = True
if failed:
raise DSLRuntimeError(
f"Incorrect {symbol.__name__} is used.",
suggestion=f"Please avoid overriding `{symbol.__name__}` from DSL package.",
)
def redirect_builtin_function(fcn):
"""
This function is used to redirect built-in function call
to the function defined in DSL package.
"""
# Only redirect if it's a built-in
if isinstance(fcn, BuiltinFunctionType) and executor._builtin_redirector:
return executor._builtin_redirector(fcn)
return fcn
python/CuTeDSL/base_dsl/ast_preprocessor.py
+376 -185
View File
File diff suppressed because it is too large Load Diff
python/CuTeDSL/base_dsl/cache_helpers.py
+1 -2
View File
@@ -139,8 +139,7 @@ def dump_cache_to_path(
dsl_name, jit_cache, cache_limit, path=default_generated_ir_path
):
log().info("JIT cache : dumping [%s] items=[%s]", dsl_name, len(jit_cache))
if not os.path.exists(path):
os.makedirs(path)
os.makedirs(path, exist_ok=True)
original_path = os.getcwd()
try:
os.chdir(path)
python/CuTeDSL/base_dsl/compiler.py
+2
View File
@@ -205,6 +205,8 @@ class CompileOptions:
self._parser.add_argument(
"--enable-device-assertions", action="store_true", default=False
)
self._parser.add_argument("--link-libraries", type=str, default="")
try:
self._options = self._parser.parse_args(options.split())
except SystemExit as e:
python/CuTeDSL/base_dsl/dsl.py
+56 -81
View File
@@ -32,13 +32,14 @@ import hashlib
from functools import lru_cache, wraps
from collections import namedtuple
from abc import ABC, abstractmethod
from typing import Any, Union, Tuple, get_origin, get_args
from types import FunctionType
from typing import Any, Union, Tuple, get_origin, get_args, List
from types import FunctionType, SimpleNamespace
import warnings
from . import typing as t
from .env_manager import EnvironmentVarManager
from .compiler import CompileOptions
from .ast_helpers import DSLOptimizationWarning
# =============================================================================
# CUDA Python
@@ -56,7 +57,7 @@ from .utils.timer import timer
from .utils.logger import setup_log, log
from .utils.stacktrace import filter_exception, walk_to_top_module, filter_stackframe
from .runtime.jit_arg_adapters import is_argument_constexpr, JitArgAdapterRegistry
from .runtime.tensor_descriptor import TensorDescriptor
from .ast_preprocessor import DSLPreprocessor
from .common import *
from .typing import (
@@ -73,12 +74,6 @@ from .._mlir import runtime as rt
from .._mlir.extras import types as T
from .._mlir.dialects import arith, math, func
# =============================================================================
# cutlass.dlpack_runtime
# =============================================================================
from .runtime.dlpack_runtime import dlpack_to_tensor_desc, mark_layout_dynamic
# =============================================================================
# Global Variables
# =============================================================================
@@ -177,6 +172,7 @@ def is_dynamic_expression(value):
return True
return False
def extract_mlir_values(obj):
"""
Given the `obj`, recursively go through it to extract all contained IR values as list of MLIR values
@@ -186,6 +182,10 @@ def extract_mlir_values(obj):
res = obj.__extract_mlir_values__()
elif isinstance(obj, (tuple, list)):
res = sum((extract_mlir_values(x) for x in obj), [])
elif isinstance(obj, SimpleNamespace):
res = []
for k, v in obj.__dict__.items():
res.extend(extract_mlir_values(v))
# Can't call is_dynamic_expression as _is_dynamic_expression depends on extract_mlir_values
elif isinstance(obj, set):
raise DSLRuntimeError(
@@ -215,6 +215,13 @@ def new_from_mlir_values(obj, values):
values = values[n_items:]
obj_ty = type(obj)
return obj_ty(res)
elif isinstance(obj, SimpleNamespace):
res = SimpleNamespace()
for k, v in obj.__dict__.items():
n_items = len(get_mlir_types(v))
res.__dict__[k] = new_from_mlir_values(v, values[:n_items])
values = values[n_items:]
return res
elif isinstance(obj, set):
raise DSLRuntimeError(
"Sets are not supported in new_from_mlir_values to ensure order preservation",
@@ -249,8 +256,6 @@ class DSLCallable:
Methods:
__call__(*args, **kwargs): Calls the wrapped function and clears it.
get_arg_spec(): Returns the argument specification of the function.
get_signature(): Returns the signature of the function.
"""
def __init__(self, func):
@@ -266,23 +271,23 @@ class DSLCallable:
assert self.func is not None, "DSLCallable is already called"
return self.func
@property
def __signature__(self):
return inspect.signature(self.__func__)
@property
def __name__(self):
return self.__func__.__name__
def get_arg_spec(self):
return inspect.getfullargspec(self.__func__)
def get_signature(self):
return inspect.signature(self.__func__)
class BaseDSL:
gpu_module = None
def __init__(
self,
*,
name: str,
dsl_package_name: List[str],
compiler_provider: Any,
pass_sm_arch_name: str,
device_compilation_only=False,
@@ -293,6 +298,7 @@ class BaseDSL:
Parameters:
- name (str): Name of DSL, used for environment variables and logging.
- package_name (str): Name of the package, used for the preprocessor.
- compiler_provider (MLIR dialect): Provider for compiler.
- pass_sm_arch_name (str): The keyword name of the SM.
- device_compilation_only (bool) : Only device code, and call it via cuda driver
@@ -330,6 +336,9 @@ class BaseDSL:
self.device_jit_decorator_name = f"@{BaseDSL.kernel.__name__}"
# set warning
if not self.envar.enable_optimization_warnings:
# By default, optimization warnings are disabled
warnings.filterwarnings("ignore", category=DSLOptimizationWarning)
if self.envar.warnings_as_errors:
warnings.filterwarnings("error")
if self.envar.warnings_ignore:
@@ -355,7 +364,7 @@ class BaseDSL:
self.compile_options = CompileOptions()
if preprocess:
self.preprocessor = DSLPreprocessor()
self.preprocessor = DSLPreprocessor(dsl_package_name)
log().info(f"Initializing {name} DSL")
log().debug(f"Logger initialized for {self.name}")
@@ -656,7 +665,7 @@ class BaseDSL:
return ir_args, ir_kwargs
@abstractmethod
def _generate_mlir_type_for_tensor_descriptor(self, tensor: TensorDescriptor):
def _generate_mlir_type_for_tensor_descriptor(self, tensor):
"""
Generate MLIR type for the tensor descriptor.
"""
@@ -671,13 +680,6 @@ class BaseDSL:
"""
pass
@abstractmethod
def _get_module_globals(self):
"""
Get the module's globals.
"""
pass
def _get_globals(self):
"""
Combines global and local variables from the current context and the
@@ -690,43 +692,21 @@ class BaseDSL:
AST preprocessor generates a new python code, so the resulting globals
dictionary is used to execute the python code.
"""
all_globals = self._get_module_globals().copy()
all_globals = {}
if self.frame:
all_globals.update(self.frame.f_globals)
all_globals.update(self.frame.f_locals)
return all_globals
@abstractmethod
def _is_tensor_descriptor(self, maybe_tensor_descriptor) -> bool:
return isinstance(
maybe_tensor_descriptor, TensorDescriptor
) or TensorDescriptor.can_transformed_to_dlpack(maybe_tensor_descriptor)
pass
@abstractmethod
def _handle_tensor_descriptor(
self, maybe_tensor, arg_name: str, need_gpu_memory: bool
) -> TensorDescriptor:
if self._is_tensor_descriptor(maybe_tensor):
tensor = (
maybe_tensor
if isinstance(maybe_tensor, TensorDescriptor)
else TensorDescriptor(maybe_tensor)
)
if need_gpu_memory and not tensor.is_in_device:
log().info(
"FAIL name=[%s] tensor=[%s] in_gpu=[%s]",
arg_name,
tensor,
tensor.is_in_device,
)
raise DSLRuntimeError(
f'Tensor "{arg_name}" is tensor "{tensor}" '
"is not in the GPU memory. "
)
return tensor
raise DSLRuntimeError(
f"Argument {arg_name} could not be transformed into a TensorDescriptor."
)
) -> Any:
pass
def _validate_arg(self, arg, arg_index, arg_name, arg_spec):
"""
@@ -882,10 +862,11 @@ class BaseDSL:
cluster: list = None
grid: list = field(default_factory=lambda: [1, 1, 1])
block: list = field(default_factory=lambda: [1, 1, 1])
smem: int = 0
smem: int = None
async_deps: list = field(default_factory=list)
has_cluster: bool = False
min_blocks_per_mp: int = 0
auto_smem: bool = False
def __post_init__(self):
if len(self.grid) != 3:
@@ -893,6 +874,10 @@ class BaseDSL:
if len(self.block) != 3:
raise DSLRuntimeError(f"Expect 3d block!")
if self.smem is None:
self.smem = 0
self.auto_smem = True
self.has_cluster = self.cluster is not None
if self.cluster is None:
self.cluster = [None, None, None]
@@ -1116,8 +1101,6 @@ class BaseDSL:
try:
result = funcBody(*ir_args, **ir_kwargs)
func.ReturnOp([])
except DSLAstPreprocessorError as pp_error:
raise pp_error
except NameError as name_error:
raise DSLRuntimeError(
f"💥💥💥 Error during runtime code generation for function `{funcBody.__name__}` 💥💥💥",
@@ -1127,11 +1110,6 @@ class BaseDSL:
except DSLRuntimeError as dsl_error:
# Throw it's already a DSL error
raise dsl_error
except Exception as general_e:
# Transform internal error to a DSL error
raise DSLRuntimeError(
f"💥💥💥 Error during runtime code generation for function `{funcBody.__name__}` 💥💥💥"
) from general_e
return module, result
# Build IR module
@@ -1328,10 +1306,8 @@ class BaseDSL:
raise DSLRuntimeError("Function body is not set.")
# Pass the actual function object to inspect.signature to get the signature.
if isinstance(self.funcBody, DSLCallable):
sig = self.funcBody.get_signature()
else:
sig = inspect.signature(self.funcBody)
sig = inspect.signature(self.funcBody)
function_name = self.funcBody.__name__
bound_args = self._get_function_bound_args(sig, function_name, *args, **kwargs)
@@ -1382,10 +1358,7 @@ class BaseDSL:
# Check the number of arguments
sig = self._check_arg_count(*args, **kwargs)
if isinstance(funcBody, DSLCallable):
args_spec = funcBody.get_arg_spec()
else:
args_spec = inspect.getfullargspec(funcBody)
args_spec = inspect.getfullargspec(funcBody)
# Canonicalize the input arguments
canonicalized_args, canonicalized_kwargs = self._canonicalize_args(
@@ -1447,7 +1420,7 @@ class BaseDSL:
return cuda_helpers.stream_create()
def _execute_cuda(
self, fname_cubin, kernel_name, grid_size, block_size, stream=None
self, fname_cubin, kernel_name, grid_size, block_size, smem_size, stream=None
):
"""
Executes a specified CUDA kernel from a cubin file, handling module loading,
@@ -1471,7 +1444,7 @@ class BaseDSL:
grid_size,
block_size,
stream,
smem_size=16000,
smem_size=smem_size,
kernel_args=self.exe_args,
)
@@ -1480,7 +1453,13 @@ class BaseDSL:
cuda_helpers.stream_sync(stream)
def _execute_by_cuda_driver(
self, kernel_generator, generate_cubin, grid_size, block_size, stream=None
self,
kernel_generator,
generate_cubin,
grid_size,
block_size,
smem_size,
stream=None,
):
"""
This function builds IR and execute the module using cuda driver.
@@ -1511,10 +1490,9 @@ class BaseDSL:
fname_cubin = generate_cubin(module, kernel_name)
# Execute a cuda kernel from cubin
if block_size is None:
# The TileIR driver should set this automatically.
block_size = self.block_size
self._execute_cuda(fname_cubin, kernel_name, grid_size, block_size, stream)
self._execute_cuda(
fname_cubin, kernel_name, grid_size, block_size, smem_size, stream
)
return ret
@@ -1587,10 +1565,7 @@ class BaseDSL:
kernelGenHelper = dkwargs.get("kernelGenHelper", None)
kernel_name = funcBody.__name__
if isinstance(funcBody, DSLCallable):
args_spec = funcBody.get_arg_spec()
else:
args_spec = inspect.getfullargspec(funcBody)
args_spec = inspect.getfullargspec(funcBody)
self.funcBody = funcBody
# Give each kernel a unique name. (The same kernel may be
python/CuTeDSL/base_dsl/env_manager.py
+21 -2
View File
@@ -58,6 +58,11 @@ def get_int_env_var(var_name, default_value=0):
return int(value) if value and value.isdigit() else default_value
@lru_cache(maxsize=None)
def has_env_var(var_name):
return os.getenv(var_name) is not None
def detect_gpu_arch(prefix):
"""
Attempts to detect the machine's GPU architecture.
@@ -256,6 +261,7 @@ class EnvironmentVarManager:
- [DSL_NAME]_ARCH: GPU architecture (default: "sm_100")
- [DSL_NAME]_WARNINGS_AS_ERRORS: Enable warnings as error (default: False)
- [DSL_NAME]_WARNINGS_IGNORE: Ignore warnings (default: False)
- [DSL_NAME]_ENABLE_OPTIMIZATION_WARNINGS: Enable warnings of optimization warnings (default: False)
- [DSL_NAME]_JIT_TIME_PROFILING: Whether or not to profile the IR generation/compilation/execution time (default: False)
- [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)
@@ -267,7 +273,6 @@ class EnvironmentVarManager:
self.prefix = prefix # change if needed
# Printing options
self.log_to_console = get_bool_env_var(f"{prefix}_LOG_TO_CONSOLE", False)
self.print_after_preprocessor = get_bool_env_var(
f"{prefix}_PRINT_AFTER_PREPROCESSOR", False
)
@@ -275,15 +280,29 @@ class EnvironmentVarManager:
self.filterStacktrace = get_bool_env_var(f"{prefix}_FILTER_STACKTRACE", True)
# File options
self.keepIR = get_bool_env_var(f"{prefix}_KEEP_IR", False)
# Logging options
self.log_to_console = get_bool_env_var(f"{prefix}_LOG_TO_CONSOLE", False)
self.log_to_file = get_bool_env_var(f"{prefix}_LOG_TO_FILE", False)
# Other options
if (
has_env_var(f"{prefix}_LOG_LEVEL")
and not self.log_to_console
and not self.log_to_file
):
log().warning(
f"Log level was set, but neither logging to file ({prefix}_LOG_TO_FILE) nor logging to console ({prefix}_LOG_TO_CONSOLE) is enabled!"
)
self.log_level = get_int_env_var(f"{prefix}_LOG_LEVEL", 1)
# Other options
self.dryrun = get_bool_env_var(f"{prefix}_DRYRUN", False)
self.arch = get_str_env_var(f"{prefix}_ARCH", detect_gpu_arch(prefix))
self.warnings_as_errors = get_bool_env_var(
f"{prefix}_WARNINGS_AS_ERRORS", False
)
self.warnings_ignore = get_bool_env_var(f"{prefix}_WARNINGS_IGNORE", False)
self.enable_optimization_warnings = get_bool_env_var(
f"{prefix}_ENABLE_OPTIMIZATION_WARNINGS", False
)
self.jitTimeProfiling = get_bool_env_var(f"{prefix}_JIT_TIME_PROFILING", False)
self.disable_file_caching = get_bool_env_var(
f"{prefix}_DISABLE_FILE_CACHING", False
python/CuTeDSL/base_dsl/runtime/__init__.py
-4
View File
@@ -14,16 +14,12 @@ This module provides a runtime utility functions that are needed for
the DSL.
"""
from . import device_tensor
from . import dlpack_types
from . import cuda
from . import tensor_descriptor
from . import jit_arg_adapters
__all__ = [
"device_tensor",
"dlpack_types",
"cuda",
"tensor_descriptor",
"jit_arg_adapters",
]
python/CuTeDSL/base_dsl/runtime/cuda.py
+1 -1
View File
@@ -309,7 +309,7 @@ def get_kernel_function(module, kernel_name):
return kernel
def launch_kernel(kernel, grid_dims, block_dims, stream, smem_size=0, kernel_args=None):
def launch_kernel(kernel, grid_dims, block_dims, stream, smem_size, kernel_args=None):
"""
Launches the CUDA kernel.
"""
python/CuTeDSL/base_dsl/runtime/tensor_descriptor.py
+7 -7
View File
@@ -183,6 +183,13 @@ class TensorDescriptor:
"""
return self.device_type == _dpack.DLDeviceType.kDLGPU
@staticmethod
def is_compatible(maybe_tensor_descriptor) -> bool:
"""Check if the object is a TensorDescriptor or can be converted to one."""
return isinstance(
maybe_tensor_descriptor, TensorDescriptor
) or TensorDescriptor.can_transformed_to_dlpack(maybe_tensor_descriptor)
def from_tensor(tensor) -> TensorDescriptor:
"""Create a TensorDescriptor from a tensor object."""
@@ -192,10 +199,3 @@ def from_tensor(tensor) -> TensorDescriptor:
def to_tensor(tensor_descriptor: TensorDescriptor):
"""Return tensor object from tensor descriptor."""
return tensor_descriptor.tensor
def is_tensor_descriptor(maybe_tensor_descriptor) -> bool:
"""Check if the object is a TensorDescriptor."""
return isinstance(
maybe_tensor_descriptor, TensorDescriptor
) or TensorDescriptor.can_transformed_to_dlpack(maybe_tensor_descriptor)
python/CuTeDSL/cutlass/cute/__init__.py
+2
View File
@@ -126,6 +126,7 @@ from .core import (
basic_copy_if,
autovec_copy,
copy,
copy_atom_call,
gemm,
# Wrapper classes
ComposedLayout,
@@ -290,6 +291,7 @@ __all__ = [
"basic_copy_if",
"autovec_copy",
"copy",
"copy_atom_call",
"gemm",
# Tensor creation
"full",
python/CuTeDSL/cutlass/cute/arch/mbar.py
+32
View File
@@ -315,3 +315,35 @@ def mbarrier_arrive(
loc=loc,
ip=ip,
)
@dsl_user_op
def cp_async_mbarrier_arrive_noinc(mbar_ptr: Pointer, *, loc=None, ip=None) -> None:
"""
Arrives on an mbarrier for async load **without incrementing** the arrival count
(`cp.async.mbarrier.arrive.shared ..., noinc=1`).
Used in the warp-specialized kernel when the non-TMA load warp(producer) is not the same
as the math/epilogue warp(consumer).
:param mbar_ptr: A pointer to the mbarrier in SMEM
:type mbar_ptr: Pointer
"""
arch = CuTeDSL._get_dsl().envar.arch
check_value_in(
arch,
[
"sm_90",
"sm_90a",
"sm_100a",
"sm_100f",
],
"arch",
)
mbar_llvm_ptr = mbar_ptr.llvm_ptr
nvvm.cp_async_mbarrier_arrive_shared(
mbar_llvm_ptr,
noinc=True,
loc=loc,
ip=ip,
)
python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py
+10 -2
View File
@@ -11,6 +11,7 @@
from functools import partial
from typing import Optional, Tuple, Union, Callable
from typing_extensions import deprecated
from cutlass.cutlass_dsl import T, dsl_user_op
@@ -642,6 +643,9 @@ def rcp_approx(a: Union[float, Float32], *, loc=None, ip=None):
@dsl_user_op
@deprecated(
"cute.arch.exp2 is deprecated, use cute.math.exp2 with `fastmath=True` instead"
)
def exp2(a: Union[float, Float32], *, loc=None, ip=None) -> Float32:
return Float32(
llvm.inline_asm(
@@ -656,15 +660,19 @@ def exp2(a: Union[float, Float32], *, loc=None, ip=None) -> Float32:
)
# TODO: add `fastmath` flag for this op
@dsl_user_op
@deprecated(
"cute.arch.exp is deprecated, use cute.math.exp with `fastmath=True` instead"
)
def exp(a: Union[float, Float32], *, loc=None, ip=None) -> Float32:
LOG2_E = 1.4426950408889634
return exp2(a * LOG2_E, loc=loc, ip=ip)
# TODO: add `fastmath` flag for this op
@dsl_user_op
@deprecated(
"cute.arch.exp_packed_f32x2 is deprecated, use cute.arch.mul_packed_f32x2 and cute.math.exp2 with `fastmath=True` instead"
)
def exp_packed_f32x2(
a: Tuple[Float32, Float32], *, loc=None, ip=None
) -> Tuple[Float32, Float32]:
python/CuTeDSL/cutlass/cute/core.py
+284 -88
View File
@@ -31,7 +31,6 @@ from typing import (
Optional,
)
from enum import Enum, auto
from typing_extensions import deprecated
from cutlass.cutlass_dsl import (
const,
@@ -1662,7 +1661,9 @@ class _Tensor(Tensor):
@dsl_user_op
def print_tensor(tensor: Tensor, *, verbose: bool = False, loc=None, ip=None):
def print_tensor(
tensor: Union[Tensor, "TensorSSA"], *, verbose: bool = False, loc=None, ip=None
):
"""Print content of the tensor in human readable format.
Outputs the tensor data in a structured format showing both metadata
@@ -1693,6 +1694,11 @@ def print_tensor(tensor: Tensor, *, verbose: bool = False, loc=None, ip=None):
[ 0.9159, 0.7577, 0.6918, 0.0754, 0.0591],
[ 0.6551, 0.1626, 0.1189, 0.0292, 0.8655]])
"""
if isinstance(tensor, TensorSSA):
tmp = make_fragment(tensor.shape, tensor.dtype)
tmp.store(tensor)
tensor = tmp
if not isinstance(tensor.type, _cute_ir.MemRefType):
raise NotImplementedError(
f"printing {tensor} is not supported because it doesn't support trivial dereferencing. "
@@ -1769,7 +1775,7 @@ def is_static(x: Union[ir.Type, ir.Value, XTuple]) -> bool:
return False
elif is_dynamic_expression(x):
return _cute_ir.is_static(x.type)
elif isinstance(x, int) or x is None:
elif isinstance(x, (bool, int, float)) or x is None:
return True
elif isinstance(x, ScaledBasis):
return x.is_static()
@@ -2241,7 +2247,7 @@ def is_weakly_congruent(
* X is a non-tuple value, OR
* X and Y are both tuples of the same rank AND all corresponding elements are weakly congruent.
Weak congruence allows scalar values to match with tuples, making it useful
Weak congruence allows scalar values to match with tuples, making it useful
for determining whether an object has a hierarchical structure "up to" another.
:param a: First object to compare
@@ -2921,33 +2927,46 @@ def flatten_to_tuple(a: Union[IntTuple, Coord, Shape, Stride]) -> tuple:
return tuple(chain.from_iterable(tuple(flatten_to_tuple(x) for x in a)))
def flatten(a: Union[IntTuple, Coord, Shape, Stride, Layout, Tensor]) -> tuple:
@overload
def flatten(a: Union[IntTuple, Coord, Shape, Stride]) -> IntTuple: ...
@overload
def flatten(a: Tensor) -> Tensor: ...
@overload
def flatten(a: Layout) -> Layout: ...
def flatten(a):
"""Flattens a CuTe data structure into a simpler form.
For tuples, this function flattens the structure into a single-level tuple.
For non-tuple types, it returns the input unchanged.
For layouts, it returns a new layout with flattened shape and stride.
For tensors, it returns a new tensor with flattened layout.
For other types, it returns the input unchanged.
:param a: The structure to flatten
:type a: Union[IntTuple, Coord, Shape, Stride, Layout, Tensor]
:return: The flattened structure
:rtype: Union[tuple, Any]
:raises NotImplementedError: If input is a Layout or Tensor
**Examples:**
.. code-block:: python
flatten((1, 2, 3)) # Returns (1, 2, 3)
flatten(((1, 2), (3, 4))) # Returns (1, 2, 3, 4)
flatten(5) # Returns 5
"""
if isinstance(a, (Layout, Tensor)):
raise NotImplementedError("flatten layout and tensor is not supported")
flatten((1, 2, 3)) # Returns (1, 2, 3)
flatten(((1, 2), (3, 4))) # Returns (1, 2, 3, 4)
flatten(5) # Returns 5
flatten(Layout(shape, stride)) # Returns Layout(flatten(shape), flatten(stride))
flatten(Tensor(layout)) # Returns Tensor(flatten(layout))
if not isinstance(a, tuple):
return a
else:
"""
if isinstance(a, Tensor):
return make_tensor(a.iterator, flatten(a.layout))
elif isinstance(a, Layout):
return make_layout(flatten(a.shape), stride=flatten(a.stride))
elif isinstance(a, tuple):
return flatten_to_tuple(a)
else:
return a
def unflatten(
@@ -4120,14 +4139,14 @@ def complement(
@dsl_user_op
def right_inverse(input: Layout, *, loc=None, ip=None) -> Layout:
if not isinstance(input, Layout):
raise TypeError(f"expects input of type Layout, but got {type(Layout)}")
raise TypeError(f"expects input of type Layout, but got {type(input)}")
return _cute_ir.right_inverse(input=input, loc=loc, ip=ip)
@dsl_user_op
def left_inverse(input: Layout, *, loc=None, ip=None) -> Layout:
if not isinstance(input, Layout):
raise TypeError(f"expects input of type Layout, but got {type(Layout)}")
raise TypeError(f"expects input of type Layout, but got {type(input)}")
return _cute_ir.left_inverse(input=input, loc=loc, ip=ip)
@@ -5156,7 +5175,6 @@ def _make_tiled_copy(atom, layout_tv, tiler_mn, *, loc=None, ip=None):
return TiledCopy(atom.op, trait)
@deprecated("Use make_tiled_copy_tv instead")
def make_tiled_copy(atom, layout_tv, tiler_mn, *, loc=None, ip=None):
"""Create a tiled type given a TV partitioner and tiler.
@@ -5434,6 +5452,14 @@ def gemm(
For MMA Atoms that require single-threaded execution, the gemm op automatically handles thread
election internally. Manual thread selection is not required in such cases.
Following dispatch rules are supported:
- Dispatch [1]: (V) x (V) => (V) => (V,1,1) x (V,1,1) => (V,1,1)
- Dispatch [2]: (M) x (N) => (M,N) => (1,M,1) x (1,N,1) => (1,M,N)
- Dispatch [3]: (M,K) x (N,K) => (M,N) => (1,M,K) x (1,N,K) => (1,M,N)
- Dispatch [4]: (V,M) x (V,N) => (V,M,N) => (V,M,1) x (V,N,1) => (V,M,N)
- Dispatch [5]: (V,M,K) x (V,N,K) => (V,M,N)
:param atom: MMA atom
:type atom: MmaAtom
:param d: Destination tensor
@@ -5454,6 +5480,27 @@ def gemm(
:rtype: None
"""
a_rank = rank(a.shape)
b_rank = rank(b.shape)
c_rank = rank(c.shape)
d_rank = rank(d.shape)
if a_rank != b_rank:
raise ValueError("`a` and `b` must have the same rank")
if c_rank != d_rank:
raise ValueError("`c` and `d` must have the same rank")
if a_rank == 1:
if c_rank > 2:
raise ValueError("`c` must have rank <= 2 when `a` has rank 1")
elif a_rank == 2:
if c_rank not in (2, 3):
raise ValueError("`c` must have rank 2 or 3 when `a` has rank 2")
elif a_rank == 3:
if c_rank != 3:
raise ValueError("`c` must have rank 3 when `a` has rank 3")
value = atom._unpack(loc=loc, ip=ip, **kwargs)
return _cute_ir.gemm(value, d.value, a.value, b.value, c.value, loc=loc, ip=ip)
@@ -5645,6 +5692,76 @@ def copy(
return _cute_ir.copy(value, src.value, dst.value, pred=pred, loc=loc, ip=ip)
@dsl_user_op
def copy_atom_call(
atom: CopyAtom,
src: Tensor,
dst: Tensor,
*,
pred: Optional[Tensor] = None,
loc=None,
ip=None,
**kwargs,
) -> None:
"""
Execute a single copy atom operation.
The copy_atom_call operation executes a copy atom with the given operands.
Following src/dst layout of atom are valid:
* ((atom_v))
* (atom_v)
Note: The format ((atom_v, rest_v)) is NOT valid for copy_atom_call since it would
require multiple atom operations, which contradicts the definition of a single copy atom call.
Examples:
.. code-block:: python
# Call a copy atom operation
cute.copy_atom_call(copy_atom, src_tensor, dst_tensor)
An additional predication tensor can be provided. If the partitioned tensors have the following
logical profile ``((ATOM_V,ATOM_REST),REST_M,...)``, the predication tensor must have a profile
consistent with ``(ATOM_REST,REST_M,...)``.
"""
if isinstance(src.type, _cute_ir.MemRefType) and isinstance(
dst.type, _cute_ir.MemRefType
):
if src.element_type.width != dst.element_type.width:
raise TypeError(
"`copy_atom_call` currently only supports equal source and destination "
"element type bit width"
)
value = atom._unpack(loc=loc, ip=ip, **kwargs)
if isinstance(pred, Tensor):
pred = pred.value
return _cute_ir.copy_atom_call(
value, src.value, dst.value, pred=pred, loc=loc, ip=ip
)
def prefetch(atom: CopyAtom, src: Tensor, *, loc=None, ip=None) -> None:
"""
The Prefetch algorithm.
The "prefetch" expects source tensors to be partitioned according to the provided Copy Atom.
Prefetch is used for loading tensors from global memory to L2.
Prefetch accepts Copy Atom but not all are allowed. Currently, only support for tma load tensor prefetch.
.. code-block:: python
cute.prefetch(tma_atom, src)
For Copy Atoms that require single-threaded execution, the copy op automatically handles thread
election internally. Manual thread selection is not required in such cases.
"""
dummy_tma_bar_ptr = make_ptr(Int64, 0, AddressSpace.smem, loc=loc, ip=ip)
value = atom._unpack(loc=loc, ip=ip, tma_bar_ptr=dummy_tma_bar_ptr)
return _cute_ir.prefetch(value, src.value, loc=loc, ip=ip)
####################################################################################################
#
# TensorSSA class (experimental)
@@ -5657,6 +5774,11 @@ class ReductionOp(Enum):
MUL = auto()
MAX = auto()
MIN = auto()
INC = auto()
DEC = auto()
AND = auto()
OR = auto()
XOR = auto()
def __str__(self):
return self.name.lower()
@@ -5697,6 +5819,7 @@ class TensorSSA(cutlass_arith.ArithValue):
self._shape = shape
self._dtype = dtype
self._layout = None
@property
def dtype(self) -> Type[Numeric]:
@@ -5776,13 +5899,26 @@ class TensorSSA(cutlass_arith.ArithValue):
):
res_type = Boolean
if lhs.shape != rhs.shape:
raise ValueError(
f"lhs and rhs must have the same shape type, but got {lhs.shape} and {rhs.shape}"
)
assert isinstance(rhs, TensorSSA), f"rhs must be TensorSSA but got {rhs}"
if not isinstance(rhs, TensorSSA):
raise TypeError(f"rhs must be TensorSSA but got {rhs}")
def _broadcast(s, t):
if s == 1:
return t
elif t == 1:
return s
elif s == t:
return s
else:
raise ValueError(f"cannot broadcast {s} and {t}")
max_rank = max(rank(lhs.shape), rank(rhs.shape))
lhs_shape = append(lhs.shape, 1, up_to_rank=max_rank)
rhs_shape = append(rhs.shape, 1, up_to_rank=max_rank)
res_shape = transform_leaf(_broadcast, lhs_shape, rhs_shape)
# broadcast to the same shape
lhs = lhs.broadcast_to(res_shape)
rhs = rhs.broadcast_to(res_shape)
if (
op in (operator.add, operator.sub)
@@ -5807,6 +5943,38 @@ class TensorSSA(cutlass_arith.ArithValue):
return res
def broadcast_to(self, target_shape: Shape, *, loc=None, ip=None) -> "TensorSSA":
"""
Broadcast the tensor to the target shape.
"""
# pad source shape to the same rank
shape = append(self.shape, 1, up_to_rank=rank(target_shape))
if shape == target_shape:
return self
def _check_broadcast(s, t):
if s != t and s != 1:
raise ValueError(
f"src_shape and target_shape must be the same when src_shape is not 1, but got {s} and {t}"
)
transform_leaf(_check_broadcast, shape, target_shape)
# 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)
# 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)
def __pow__(self, other, *, loc=None, ip=None) -> "TensorSSA":
"""
Returns the results of tensor^other.
@@ -6093,6 +6261,16 @@ class TensorSSA(cutlass_arith.ArithValue):
"""
return self._apply_op(operator.and_, other, flip=True, loc=loc, ip=ip)
def __neg__(self, *, loc=None, ip=None) -> "TensorSSA":
"""
Returns the negation of the tensor.
:return: The element-wise negation of the tensor
:rtype: TensorSSA
"""
return self._apply_op(operator.sub, 0, flip=True, loc=loc, ip=ip)
def _flatten_shape_and_coord(self, crd, *, loc=None, ip=None):
# Coalesce and flatten source layout at terminal of coordinate
# (N_0,(N_1,...), ...) -> (N_0,N_1,N_2,...)
@@ -6158,17 +6336,13 @@ class TensorSSA(cutlass_arith.ArithValue):
if crd is None:
return self
if not has_underscore(crd) or depth(crd) == 0:
idx = crd2idx(crd, make_layout(self._shape))
if is_static(idx):
res = vector.extract(
self, dynamic_position=[], static_position=[idx], loc=loc, ip=ip
)
else:
res = vector.extract(
self, dynamic_position=[crd], static_position=[], loc=loc, ip=ip
)
return self.dtype(res)
if not has_underscore(crd):
if self._layout is None:
self._layout = make_layout(self._shape, loc=loc, ip=ip)
idx = crd2idx(crd, self._layout, loc=loc, ip=ip)
idx_val = as_numeric(idx).ir_value(loc=loc, ip=ip)
res_val = vector.extractelement(self, position=idx_val, loc=loc, ip=ip)
return self.dtype(res_val)
if not is_static(crd):
raise ValueError("dynamic coordinate is not supported")
@@ -6274,7 +6448,7 @@ class TensorSSA(cutlass_arith.ArithValue):
:type op: operator
:param init_val: The initial value for the reduction
:type init_val: numeric
:param reduction_profile: Specifies which dimensions to reduce. Dimensions marked with '_' are kept.
:param reduction_profile: Specifies which dimensions to reduce. Dimensions marked with `None` are kept.
:type reduction_profile: Coord
:return: The reduced tensor
@@ -6289,9 +6463,9 @@ class TensorSSA(cutlass_arith.ArithValue):
reduce(f32 o (4, 5))
=> f32
reduce(f32 o (4, (5, 4)), reduction_profile=(_, 1))
reduce(f32 o (4, (5, 4)), reduction_profile=(None, 1))
=> f32 o (4,)
reduce(f32 o (4, (5, 4)), reduction_profile=(_, (_, 1)))
reduce(f32 o (4, (5, 4)), reduction_profile=(None, (None, 1)))
=> f32 o (4, (5,))
"""
# short-cut to no-op
@@ -6354,21 +6528,6 @@ class TensorSSA(cutlass_arith.ArithValue):
return self._build_result(res_vect, res_shp, loc=loc, ip=ip)
def _get_attr_for_type(ty, value):
if isinstance(ty, ir.IntegerType):
return ir.IntegerAttr.get(ty, value.to(int))
elif isinstance(ty, ir.FloatType):
return ir.FloatAttr.get(ty, value.to(float))
else:
raise TypeError(f"unsupported type: {ty}")
def _splat(res_ty, fill_value):
elem_attr = _get_attr_for_type(res_ty.element_type, fill_value)
vect_attr = ir.DenseElementsAttr.get_splat(res_ty, elem_attr)
return arith.constant(res_ty, vect_attr)
@dsl_user_op
def full(shape, fill_value, dtype: Type[Numeric], *, loc=None, ip=None) -> TensorSSA:
"""
@@ -6389,9 +6548,14 @@ def full(shape, fill_value, dtype: Type[Numeric], *, loc=None, ip=None) -> Tenso
if isinstance(fill_value, (ir.Value, int, float, bool)):
fill_value = dtype(fill_value)
elif isinstance(fill_value, Numeric):
fill_value = fill_value.to(dtype, loc=loc, ip=ip)
else:
raise ValueError(f"Expected fill_value be numeric type, but got {fill_value}")
res_mlir_type = T.vector(size, dtype.mlir_type)
return TensorSSA(_splat(res_mlir_type, fill_value), shape, dtype)
res_ty = T.vector(size, dtype.mlir_type)
res_val = vector.splat(res_ty, fill_value.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
return TensorSSA(res_val, shape, dtype)
def full_like(
@@ -6547,7 +6711,7 @@ class struct:
**Usage:**
.. code-block::
.. code-block:: python
# Supports base_dsl scalar int/float elements, array and nested struct:
@cute.struct
@@ -6661,7 +6825,8 @@ class struct:
Initializes a new memory range.
:param dtype: The data type.
:param size: The size of the memory range in bytes.
:param size: Size of the memory range in bytes. A size of **0** is accepted, but in that
case the range can only be used for its address (e.g. as a partition marker).
:param base: The base address of the memory range.
"""
self._dtype = dtype
@@ -6673,9 +6838,9 @@ class struct:
Returns start pointer to the data in this memory range.
:return: A pointer to the start of the memory range.
:raises AssertionError: If the size of the memory range is not greater than zero.
:raises AssertionError: If the size of the memory range is negative.
"""
assert self._size > 0
assert self._size >= 0
return recast_ptr(self._base, dtype=self._dtype)
def get_tensor(self, layout, swizzle=None, dtype=None):
@@ -6716,31 +6881,48 @@ class struct:
:param v: The object to align. Must be a struct, MemRange, or a scalar type.
:param align: The alignment value to set.
:return: A copy of the object with the specified alignment.
:raises TypeError: If the object is not a struct, MemRange, or a scalar type.
:ivar _dtype: The data type to be aligned.
:ivar _align: The alignment of the data type.
"""
_dtype = None
_align = None
def __new__(cls, name, bases, dct):
return super().__new__(cls, name, bases, dct)
def __getitem__(cls, params) -> Any:
if len(params) == 2:
obj, align = params
dtype, align = params
assert align > 0
else:
raise TypeError("Invalid struct.Align Arguments")
# make a copy of type and mark alignment
if struct._is_scalar_type(obj) or isinstance(
obj, (struct, struct._MemRangeMeta)
if not struct._is_scalar_type(dtype) and not isinstance(
dtype, (struct, struct._MemRangeMeta)
):
new_obj = py_copy.copy(obj)
setattr(new_obj, "_struct_alignment_", align)
return new_obj
else:
raise TypeError(
"align only can be applied to sturct/MemRange/base_dsl scalar"
"align only can be applied to struct/MemRange/base_dsl scalar"
)
# Create new class with alignment
new_cls = type(
f"struct.Align[{dtype.__name__}, {align}]",
(struct.Align,),
{"_dtype": dtype, "_align": align},
)
return new_cls
@property
def dtype(cls):
return cls._dtype
@property
def align(cls):
return cls._align
class Align(metaclass=_AlignMeta):
"""
Aligns the given type by `Align[T, alignment]`.
@@ -6768,6 +6950,7 @@ class struct:
:raises TypeError: If the struct is empty.
"""
self._cls = cls
self.__name__ = f"struct::{cls.__name__}"
# Get the class annotations
self._annotations = cls.__annotations__
# Create a dictionary to store the offsets
@@ -6780,12 +6963,10 @@ class struct:
raise TypeError("Empty struct is not supported!")
for name, object in self._annotations.items():
# get alignment of object
def alignof(object, default: int = 1):
return getattr(object, "_struct_alignment_", default)
# alignment for the next offset
def align_offset(offset, align):
return (offset + (align - 1)) & ~(align - 1)
sub_align = 1
if isinstance(object, struct._AlignMeta):
sub_align = object.align
object = object.dtype
# switch addition order to support dynamic size
def add_offset(val):
@@ -6793,35 +6974,37 @@ class struct:
# size of scalar
if struct._is_scalar_type(object):
dtype_size = object.width // 8
sub_align = alignof(object, dtype_size)
offset = align_offset(offset, sub_align)
dtype_size = max(1, object.width // 8)
sub_align = max(dtype_size, sub_align)
offset = self.align_offset(offset, sub_align)
self._offsets[name] = offset
offset = add_offset(dtype_size)
# size of array is size_in_bytes, alignment is elem_size
elif isinstance(object, struct._MemRangeMeta):
if object.size == 0:
continue # skip empty array
sub_align = alignof(object, max(1, object.elem_width // 8))
offset = align_offset(offset, sub_align)
# Allow empty array as a free marker-only struct member.
# Use max(sub_align, ) because we might have in the future some
# object.elem_width less than 8, such as fp4, bit and others,
# and align_offset() does not support an alignment of 0.
sub_align = max(object.elem_width // 8, sub_align)
offset = self.align_offset(offset, sub_align)
self._offsets[name] = offset
offset = add_offset(object.size_in_bytes)
# size of struct
elif isinstance(object, struct):
sub_align = max(object.__alignof__(), alignof(object))
offset = align_offset(offset, sub_align)
sub_align = max(object.__alignof__(), sub_align)
offset = self.align_offset(offset, sub_align)
self._offsets[name] = offset
offset = add_offset(object.__sizeof__())
else:
raise TypeError(
f"Struct element only support sturct/array/base_dsl scalar, "
f"Struct element only support struct/array/base_dsl scalar, "
f"but got {object}"
)
# Total aligment determined by the strictest requirement
alignment = max(alignment, sub_align)
# Total size determined by alignment
self._align_of = alignment
self._size_of = align_offset(offset, alignment)
self._size_of = self.align_offset(offset, alignment)
# create the __init__ method for decorated struct
def __call__(self, base: Any) -> None:
@@ -6840,6 +7023,8 @@ class struct:
setattr(cls, "_base", base)
for name, off in self._offsets.items():
obj = self._annotations[name]
if isinstance(obj, struct._AlignMeta):
obj = obj.dtype
if struct._is_scalar_type(obj):
new_obj = recast_ptr(base + off, dtype=obj)
setattr(cls, name, new_obj)
@@ -6851,7 +7036,7 @@ class struct:
setattr(cls, name, new_obj)
else:
raise TypeError(
f"Struct element only support sturct/array/base_dsl scalar, "
f"Struct element only support struct/array/base_dsl scalar, "
f"but got {obj}"
)
return cls
@@ -6872,3 +7057,14 @@ class struct:
# get alignment
def __alignof__(self) -> int:
return self._align_of
# util func for aligning offset
@staticmethod
def align_offset(offset, align):
"""
Return the round-up offset up to the next multiple of align.
"""
assert align > 0 and not (
align & (align - 1)
), "align should be a strictly positive power of 2."
return (offset + (align - 1)) & ~(align - 1)
python/CuTeDSL/cutlass/cute/math.py
+154 -63
View File
@@ -10,16 +10,53 @@
# is strictly prohibited.
from .core import TensorSSA
from .typing import Numeric
from cutlass._mlir.dialects import math, arith
from typing import Callable, Union
def acos(a: TensorSSA) -> TensorSSA:
def _math_op(func: Callable, fastmath: bool, *args, **kwargs):
"""Dispatch the function to either a TensorSSA or a Numeric(Float).
:param func: The function to dispatch
:param args: The input tensor or scalar
:param kwargs: The input tensor or scalar
"""
arg_type = type(args[0])
for arg in args:
if not isinstance(arg, TensorSSA) and (
not isinstance(arg, Numeric) or not type(arg).is_float
):
raise TypeError(
f"Expected a TensorSSA or Numeric(Float), but got {type(arg)}"
)
if not isinstance(arg, arg_type):
raise TypeError(
f"Expected all inputs to be of type {arg_type}, but got {type(arg)}"
)
fastmath_flag = arith.FastMathFlags.fast if fastmath else arith.FastMathFlags.none
if isinstance(args[0], TensorSSA):
return TensorSSA(
func(*args, fastmath=fastmath_flag), args[0].shape, args[0].dtype
)
else:
args = [a.ir_value() for a in args]
return func(*args, fastmath=fastmath_flag)
def acos(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise arc cosine of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the arc cosine of each element in input tensor
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -29,16 +66,20 @@ def acos(a: TensorSSA) -> TensorSSA:
y = x.load() # Load values
z = acos(y) # Compute arc cosine
"""
return TensorSSA(math.acos(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.acos, fastmath, a)
def asin(a: TensorSSA) -> TensorSSA:
def asin(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise arc sine of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the arc sine of each element in input tensor
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -48,18 +89,20 @@ def asin(a: TensorSSA) -> TensorSSA:
y = x.load() # Load values
z = asin(y) # Compute arc sine
"""
return TensorSSA(math.asin(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.asin, fastmath, a)
def atan(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def atan(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise arc tangent of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the arc tangent of each element in input tensor
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -70,23 +113,25 @@ def atan(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
z = atan(y) # Compute arc tangent
"""
raise NotImplementedError("atan is not implemented")
return TensorSSA(math.atan(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.atan, fastmath, a)
def atan2(a: TensorSSA, b: TensorSSA, fastmath: bool = False) -> TensorSSA:
def atan2(
a: Union[TensorSSA, Numeric], b: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise arc tangent of two tensors.
Computes atan2(a, b) element-wise. The function atan2(a, b) is the angle in radians
between the positive x-axis and the point given by the coordinates (b, a).
:param a: First input tensor (y-coordinates)
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param b: Second input tensor (x-coordinates)
:type b: TensorSSA
:type b: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the arc tangent of a/b element-wise
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -96,20 +141,20 @@ def atan2(a: TensorSSA, b: TensorSSA, fastmath: bool = False) -> TensorSSA:
x = cute.make_fragment(ptr2, layout).load() # x coordinates
theta = atan2(y, x) # Compute angles
"""
return TensorSSA(
math.atan2(a, b, fastmath=arith.FastMathFlags.none), a.shape, a.dtype
)
return _math_op(math.atan2, fastmath, a, b)
def cos(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def cos(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise cosine of the input tensor.
:param a: Input tensor (in radians)
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the cosine of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -119,21 +164,23 @@ def cos(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = cos(y) # Compute cosine
"""
return TensorSSA(math.cos(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.cos, fastmath, a)
def erf(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def erf(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise error function of the input tensor.
The error function is defined as:
erf(x) = 2/√π ∫[0 to x] exp(-t²) dt
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the error function value for each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -143,18 +190,43 @@ def erf(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = erf(y) # Compute error function
"""
return TensorSSA(math.erf(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.erf, fastmath, a)
def exp2(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def exp(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise exponential of the input tensor.
:param a: Input tensor
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the exponential of each element
:rtype: Union[TensorSSA, Numeric]
Example:
.. code-block::
x = cute.make_fragment(layout) # Create tensor
y = x.load() # Load values
z = exp(y) # Compute exponential
"""
return _math_op(math.exp, fastmath, a)
def exp2(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise base-2 exponential of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing 2 raised to the power of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -164,18 +236,20 @@ def exp2(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = exp2(y) # Compute 2^x
"""
return TensorSSA(math.exp2(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.exp2, fastmath, a)
def log(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def log(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise natural logarithm of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the natural logarithm of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -185,18 +259,20 @@ def log(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = log(y) # Compute natural logarithm
"""
return TensorSSA(math.log(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.log, fastmath, a)
def log2(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def log2(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise base-2 logarithm of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the base-2 logarithm of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -206,18 +282,20 @@ def log2(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = log2(y) # Compute log base 2
"""
return TensorSSA(math.log2(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.log2, fastmath, a)
def log10(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def log10(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise base-10 logarithm of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the base-10 logarithm of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -227,20 +305,22 @@ def log10(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = log10(y) # Compute log base 10
"""
return TensorSSA(math.log10(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.log10, fastmath, a)
def rsqrt(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def rsqrt(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise reciprocal square root of the input tensor.
Computes 1/√x element-wise.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the reciprocal square root of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -250,18 +330,20 @@ def rsqrt(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = rsqrt(y) # Compute 1/√x
"""
return TensorSSA(math.rsqrt(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.rsqrt, fastmath, a)
def sin(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def sin(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise sine of the input tensor.
:param a: Input tensor (in radians)
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the sine of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -271,18 +353,20 @@ def sin(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = sin(y) # Compute sine
"""
return TensorSSA(math.sin(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.sin, fastmath, a)
def sqrt(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def sqrt(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise square root of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the square root of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -292,16 +376,20 @@ def sqrt(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = sqrt(y) # Compute square root
"""
return TensorSSA(math.sqrt(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.sqrt, fastmath, a)
def tan(a: TensorSSA) -> TensorSSA:
def tan(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise tangent of the input tensor.
:param a: Input tensor (in radians)
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the tangent of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -311,18 +399,20 @@ def tan(a: TensorSSA) -> TensorSSA:
y = x.load() # Load values
z = tan(y) # Compute tangent
"""
return TensorSSA(math.tan(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.tan, fastmath, a)
def tanh(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def tanh(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise hyperbolic tangent of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the hyperbolic tangent of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -332,7 +422,7 @@ def tanh(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = tanh(y) # Compute hyperbolic tangent
"""
return TensorSSA(math.tanh(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.tanh, fastmath, a)
__all__ = [
@@ -342,6 +432,7 @@ __all__ = [
"atan2",
"cos",
"erf",
"exp",
"exp2",
"log",
"log10",
python/CuTeDSL/cutlass/cute/nvgpu/common.py
+43 -2
View File
@@ -8,7 +8,7 @@
# 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 dataclasses import dataclass
from typing import Type, Optional
@@ -101,6 +101,42 @@ class MmaUniversalTrait(core.Trait):
####################################################################################################
class MemoryOrder(enum.Enum):
WEAK = _cute_ir.MemOrderKind.WEAK
RELAXED = _cute_ir.MemOrderKind.RELAXED
ACQUIRE = _cute_ir.MemOrderKind.ACQUIRE
RELEASE = _cute_ir.MemOrderKind.RELEASE
ACQ_REL = _cute_ir.MemOrderKind.ACQ_REL
SC = _cute_ir.MemOrderKind.SC
MMIO = _cute_ir.MemOrderKind.MMIO
CONSTANT = _cute_ir.MemOrderKind.CONSTANT
VOLATILE = _cute_ir.MemOrderKind.VOLATILE
def __str__(self) -> str:
return f"{self.__class__.__name__}.{self.name}"
def __repr__(self) -> str:
return f"<{self.__class__.__name__}.{self.name}>"
def _to_ir(self) -> _cute_ir.MemOrderKind:
return self.value
class MemoryScope(enum.Enum):
CTA = _cute_ir.MemScopeKind.CTA
CLUSTER = _cute_ir.MemScopeKind.CLUSTER
GPU = _cute_ir.MemScopeKind.GPU
SYS = _cute_ir.MemScopeKind.SYS
def __str__(self) -> str:
return f"{self.__class__.__name__}.{self.name}"
def __repr__(self) -> str:
return f"<{self.__class__.__name__}.{self.name}>"
def _to_ir(self) -> _cute_ir.MemScopeKind:
return self.value
@dataclass(frozen=True)
class CopyUniversalOp(core.CopyOp):
"""
@@ -133,13 +169,18 @@ class CopyUniversalOp(core.CopyOp):
**kwargs,
) -> "CopyUniversalTrait":
num_bits_per_copy = kwargs.get("num_bits_per_copy", 0)
memory_order = kwargs.get("memory_order", MemoryOrder.WEAK)
memory_scope = kwargs.get("memory_scope", MemoryScope.CTA)
if not isinstance(num_bits_per_copy, int) or (num_bits_per_copy < 0):
raise ValueError(
"expects a 'num_bits_per_copy' kw argument of type int that is non-negative "
f"when creating a copy Atom for {self.__class__.__name__}"
)
ty = _cute_nvgpu_ir.CopyAtomSIMTSyncCopyType.get(
copy_internal_type.mlir_type, num_bits_per_copy
copy_internal_type.mlir_type,
num_bits_per_copy,
memory_order._to_ir(),
memory_scope._to_ir(),
)
return CopyUniversalTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
python/CuTeDSL/cutlass/cute/nvgpu/cpasync/__init__.py
+1
View File
@@ -23,6 +23,7 @@ __all__ = [
"CopyBulkTensorTileG2SOp",
"CopyBulkTensorTileG2SMulticastOp",
"CopyBulkTensorTileS2GOp",
"CopyReduceBulkTensorTileS2GOp",
#
# helpers.py
#

Some files were not shown because too many files have changed in this diff Show More