update 3.8 v2 (#2112)

* update 3.8 v2

* update 3.8

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
This commit is contained in:
Yujia Zhai
2025-02-19 22:03:14 -05:00
committed by GitHub
co-authored by yuzhai
parent e9627ce55b
commit b84e9802d8
166 changed files with 3986 additions and 4037 deletions
@@ -121,32 +121,25 @@ using StageCountType = cutlass::gemm::collective::StageCountAuto; // S
// Runtime Cluster Shape
using ClusterShape = Shape<int32_t,int32_t,_1>;
// For Static Cluster Shape:
// using ClusterShape = Shape<_2,_1,_1>; // for example
// using AtomThrShape = decltype(shape_div(ClusterShape{}, Shape<_2,_1,_1>{})); // for 2SM config
// using OutputTileShape = decltype(shape_div(ClusterTileShape{}, ClusterShape{})); // for epilogue builder
// using MmaTileShape = decltype(shape_div(ClusterTileShape{}, AtomThrShape{})); // for mainloop builder
// Different configs for 1SM and 2SM MMA kernel
struct MMA1SMConfig {
using MmaTileShape = Shape<_128,_256,Int<128 / sizeof(ElementA)>>;
using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmSm100; // Kernel to launch
using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm; // Epilogue to launch
using OutputTileShape = decltype(shape_div(MmaTileShape{}, Shape<_1,_1,_1>{}));
};
struct MMA2SMConfig {
using MmaTileShape = Shape<_256,_256,Int<128 / sizeof(ElementA)>>;
using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmSm100; // Kernel to launch
using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized2Sm; // Epilogue to launch
using OutputTileShape = decltype(shape_div(MmaTileShape{}, Shape<_2,_1,_1>{}));
};
template <typename ScheduleConfig>
struct GivenGemmSchedule {
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
ArchTag, OperatorClass,
typename ScheduleConfig::OutputTileShape, ClusterShape,
typename ScheduleConfig::MmaTileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementAccumulator,
ElementC, LayoutC *, AlignmentC,
@@ -143,31 +143,23 @@ using StageCountType = cutlass::gemm::collective::StageCountAuto; // S
// Runtime Cluster Shape
using ClusterShape = Shape<int32_t,int32_t,_1>;
/* // For Static Cluster Shape:
use ClusterShape = Shape<_2,_1,_1> for example
using AtomThrShape = decltype(shape_div(ClusterShape{}, Shape<_2,_1,_1>{})); // for 2SM config
using OutputTileShape = decltype(shape_div(ClusterTileShape{}, ClusterShape{})); // for epilogue builder
using MmaTileShape = decltype(shape_div(ClusterTileShape{}, AtomThrShape{})); // for mainloop builder
*/
// Different configs for 1SM and 2SM MMA kernel
struct MMA1SMConfig {
using MmaTileShape = Shape<_128,_256,_256>;
using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmNvf4Sm100; // Kernel to launch
using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm; // Epilogue to launch
using OutputTileShape = decltype(shape_div(MmaTileShape{}, Shape<_1,_1,_1>{}));
};
struct MMA2SMConfig {
using MmaTileShape = Shape<_256,_256,_256>;
using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmNvf4Sm100; // Kernel to launch
using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized2Sm; // Epilogue to launch
using OutputTileShape = decltype(shape_div(MmaTileShape{}, Shape<_2,_1,_1>{}));
};
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
ArchTag, EpilogueOperatorClass,
typename MMA1SMConfig::OutputTileShape, ClusterShape,
typename MMA1SMConfig::MmaTileShape, ClusterShape,
Shape<_128,_64>,
ElementAccumulator, ElementAccumulator,
ElementC, LayoutC *, AlignmentC,
@@ -195,7 +187,7 @@ using Gemm = Gemm1SM;
using CollectiveEpilogue2SM = typename cutlass::epilogue::collective::CollectiveBuilder<
ArchTag, EpilogueOperatorClass,
typename MMA2SMConfig::OutputTileShape, ClusterShape,
typename MMA2SMConfig::MmaTileShape, ClusterShape,
Shape<_128,_64>,
ElementAccumulator, ElementAccumulator,
ElementC, LayoutC *, AlignmentC,
@@ -49,7 +49,7 @@ set(TEST_SMALL_LARGE_GROUP --m=128 --n=128 --groups=50 --iterations=0)
set(TEST_RANDOM_PERF --iterations=10) # Random problem sizes
set(TEST_RANDOM_PERF_LARGE_GROUP --groups=50 --iterations=10) # Random problem sizes
if(NOT CUTLASS_NVCC_ARCHS STREQUAL "100")
if (CUTLASS_NVCC_ARCHS MATCHES 100a)
cutlass_example_add_executable(
75_blackwell_grouped_gemm
75_blackwell_grouped_gemm.cu