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

Updates for CUTLASS 3.5.0 (#1468)

Browse Source
This commit is contained in:
Vijay Thakkar
2024-04-11 21:33:40 -04:00
committed by GitHub
parent a40e08e9d5
commit 7d49e6c7e2
171 changed files with 7526 additions and 1888 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/gemm/device/gemm_testbed_3x.hpp
-1
View File
@@ -1336,7 +1336,6 @@ struct TestbedImpl {
{
using namespace cute;
auto problem_shape_MNKL = cute::append<4>(problem_size, 1);
auto mainloop_params = collective_mma_inputs.to_host_args(problem_size);
auto epilogue_params = collective_epilogue.to_host_args(problem_size);
test/unit/gemm/device/sm90_gemm_f32_f32_f32_tensor_op_f32.cu
+44
View File
@@ -163,6 +163,50 @@ TEST(SM90_Device_Gemm_f32t_f32t_f32n_tensor_op_gmma_f32, 128x128x32_1x1x1_cooper
EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
}
TEST(SM90_Device_Gemm_f32t_f32t_f32n_tensor_op_gmma_f32, 128x128x32_1x1x1_cooperative_narrow_wgmma) {
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
Shape<_128,_128,_32>, Shape<_1,_1,_1>,
cutlass::epilogue::collective::EpilogueTileAuto,
float, float,
float, LayoutC, 4,
float, LayoutC, 4,
cutlass::epilogue::TmaWarpSpecializedCooperative
>::CollectiveOp;
// Manually configure a half-tile wide MMA instruction
using CollectiveMainloop = cutlass::gemm::collective::CollectiveMma<
cutlass::gemm::MainloopSm90TmaGmmaWarpSpecialized<5, Shape<_1,_1,_1>, cutlass::gemm::KernelTmaWarpSpecializedCooperative>,
Shape<_128,_128,_32>,
float,
cutlass::detail::TagToStrideA_t<LayoutA>,
float,
cutlass::detail::TagToStrideB_t<LayoutB>,
decltype(cute::make_tiled_mma(cute::SM90_64x64x8_F32TF32TF32_SS_TN{}, Layout<Shape<_2,_1,_1>>{})),
cute::SM90_TMA_LOAD,
cute::GMMA::Layout_K_SW128_Atom<tfloat32_t>,
void,
cute::identity,
cute::SM90_TMA_LOAD,
cute::GMMA::Layout_K_SW128_Atom<tfloat32_t>,
void,
cute::identity
>;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue
>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
}
///////////////////////////////////////////////////////////////////////////////
#endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
test/unit/gemm/device/sm90_gemm_f8_f8_f8_tensor_op_fp32.cu
+10 -1
View File
@@ -54,6 +54,7 @@
#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
using namespace cute;
///////////////////////////////////////////////////////////////////////////////
//////////////////////////////// output: E4M3 /////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
@@ -760,7 +761,8 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_e5m2n_tensor_op_gmma_f32, 64x128x128_2x4x1_non
EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>());
}
// Use Hopper FP8+AUX from 12.1
#if (!((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ == 0)))
///////////////////////////////////////////////////////////////////////////////
///////////////////////// output: E4M3 + Aux Tensor ///////////////////////////
@@ -808,6 +810,7 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_e4m3n_tensor_op_gmma_f32, 64x128x128_aux_tenso
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>());
}
#endif
///////////////////////////////////////////////////////////////////////////////
////////////////////////////////// FP8 Accum /////////////////////////////////
@@ -990,6 +993,10 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_e4m3n_tensor_op_gmma_f32, 64x128x128_bias_bf16
EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>());
}
// Use Hopper FP8+AUX from 12.1
#if (!((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ == 0)))
///////////////////////////////////////////////////////////////////////////////
///////////////////// output: E4M3 + Aux Tensor + Bias/////////////////////////
///////////////////////////////////////////////////////////////////////////////
@@ -1142,6 +1149,8 @@ TEST(SM90_Device_Gemm_e4m3t_e5m2n_e4m3n_tensor_op_gmma_f32, 64x128x128_aux_tenso
EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>());
}
#endif
///////////////////////////////////////////////////////////////////////////////
//////////////////////////////// TMA epilogue /////////////////////////////////
///////////////////////////////////////////////////////////////////////////////