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

cutlass 3.9 update (#2255)

Browse Source 
* cutlass 3.9 update

* rebase

* fixes out of shared memory for blockwise Blackwell

* doc format

* fix issue 2253

* disable host ref by default

* fix sm120 smem capacity

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
Yujia Zhai
2025-04-24 12:42:40 -07:00
committed by GitHub
parent 8e345c5c5b
commit 331a1f5b3f
143 changed files with 18089 additions and 5935 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
test/unit/cute/core/tuple.cpp
View File

@@ -510,7 +510,7 @@ void test_sizes_and_not_storing_empty_types() {
} // namespace test
TEST(CuTe_core, PackedTuple2)
TEST(CuTe_core, PackedTuple)
{
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("tuple");
@@ -522,7 +522,7 @@ TEST(CuTe_core, PackedTuple2)
pt_test::test_sizes_and_not_storing_empty_types();
}
TEST(CuTe_core, PackedTuple2Get) {
TEST(CuTe_core, PackedTupleGet) {
using cute::tuple;
using pt_test::Empty;
using pt_test::Nonempty;
@@ -678,6 +678,42 @@ TEST(CuTe_core, PackedTuple2Get) {
}
}
TEST(CuTe_core, PackedTupleGetValueCategory) {
using cute::tuple;
using pt_test::Empty;
using pt_test::Nonempty;
tuple<Nonempty<int>, int, Empty<42>> tup(Nonempty<int>{42}, 7, Empty<42>{});
// Lvalue ref
decltype(auto) t0 = cute::get<0>(tup);
decltype(auto) t1 = cute::get<1>(tup);
decltype(auto) t2 = cute::get<2>(tup);
EXPECT_TRUE((cute::is_same_v<decltype(t0), Nonempty<int>&>));
EXPECT_TRUE((cute::is_same_v<decltype(t1), int&>));
EXPECT_TRUE((cute::is_same_v<decltype(t2), Empty<42>>));
// Const lvalue ref
auto const& ctup = tup;
decltype(auto) ct0 = cute::get<0>(ctup);
decltype(auto) ct1 = cute::get<1>(ctup);
decltype(auto) ct2 = cute::get<2>(ctup);
EXPECT_TRUE((cute::is_same_v<decltype(ct0), Nonempty<int> const&>));
EXPECT_TRUE((cute::is_same_v<decltype(ct1), int const&>));
EXPECT_TRUE((cute::is_same_v<decltype(ct2), Empty<42>>));
// Rvalue ref
decltype(auto) r0 = cute::get<0>(cute::move(tup));
decltype(auto) r1 = cute::get<1>(cute::move(tup));
decltype(auto) r2 = cute::get<2>(cute::move(tup));
EXPECT_TRUE((cute::is_same_v<decltype(r0), Nonempty<int>&&>));
EXPECT_TRUE((cute::is_same_v<decltype(r1), int&&>));
EXPECT_TRUE((cute::is_same_v<decltype(r2), Empty<42>>));
}
namespace pt_test {
// An empty class type to which Empty is convertible.
@@ -705,14 +741,14 @@ TEST(CuTe_core, PackedTupleConstexprDefaultConstruction) {
using pt_test::Empty;
{
[[maybe_unused]] constexpr cute::detail::ESO_t<Empty<0>> eso1{};
[[maybe_unused]] constexpr cute::detail::ESO_t<int64_t> eso2{};
[[maybe_unused]] constexpr cute::eso::ESO_t<Empty<0>> eso1{};
[[maybe_unused]] constexpr cute::eso::ESO_t<int64_t> eso2{};
}
{
[[maybe_unused]] constexpr cute::detail::ESO_t<Empty<0>, Empty<1>> eso0{};
[[maybe_unused]] constexpr cute::detail::ESO_t<int64_t, Empty<1>> eso1{};
[[maybe_unused]] constexpr cute::detail::ESO_t<Empty<0>, int64_t> eso2{};
[[maybe_unused]] constexpr cute::detail::ESO_t<int64_t, int64_t> eso3{};
[[maybe_unused]] constexpr cute::eso::ESO_t<Empty<0>, Empty<1>> eso0{};
[[maybe_unused]] constexpr cute::eso::ESO_t<int64_t, Empty<1>> eso1{};
[[maybe_unused]] constexpr cute::eso::ESO_t<Empty<0>, int64_t> eso2{};
[[maybe_unused]] constexpr cute::eso::ESO_t<int64_t, int64_t> eso3{};
}
}

4
test/unit/gemm/device/gemm_splitk_serial_tensor_op_sm75.cu
View File

@@ -92,7 +92,7 @@ TEST(SM75_Device_GemmSplitKSerial_f16n_f16n_f16t_tensor_op_f32, 128x256x32_64x64
cutlass::gemm::GemmShape<128, 256, 32>,
cutlass::gemm::GemmShape<64, 64, 32>,
cutlass::gemm::GemmShape<16, 8, 8>,
cutlass::epilogue::thread::LinearCombinationRelu<
cutlass::epilogue::thread::LinearCombination<
ElementOutput,
128 / cutlass::sizeof_bits<ElementOutput>::value,
ElementAccumulator,
@@ -105,7 +105,7 @@ TEST(SM75_Device_GemmSplitKSerial_f16n_f16n_f16t_tensor_op_f32, 128x256x32_64x64
kSplitKSerial
>;
bool result = test::gemm::device::TestAllGemm<Gemm, true>();
bool result = test::gemm::device::TestAllGemm<Gemm, false>();
EXPECT_TRUE(result);
}

60
test/unit/gemm/device/gemm_testbed_3x_ptr_array.hpp
View File

@@ -365,7 +365,6 @@ struct HostCollectiveMainloop {
//
// Allocate the GEMM workspace
//
// for pointer array problem_shapes.groups() is 1
tensors_A.clear();
@@ -551,7 +550,7 @@ struct HostCollectiveMainloop<cutlass::gemm::KernelPtrArrayTmaWarpSpecializedBlo
static constexpr int SFVecSize = Gemm::GemmKernel::CollectiveMainloop::SFVecSize;
using ElementSF = typename Gemm::GemmKernel::ElementSF;
using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF;
using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
using Blk_MN = typename Sm1xxBlkScaledConfig::Blk_MN;
using Blk_SF = typename Sm1xxBlkScaledConfig::Blk_SF;
@@ -619,6 +618,7 @@ struct HostCollectiveMainloop<cutlass::gemm::KernelPtrArrayTmaWarpSpecializedBlo
//
// Allocate the GEMM workspace
//
tensors_A.clear();
tensors_B.clear();
stride_a_host.clear();
@@ -786,6 +786,56 @@ struct HostCollectiveMainloop<cutlass::gemm::KernelPtrArrayTmaWarpSpecializedBlo
}
};
//
// Block Scaled Gemm Input Operands : A , B, scalefactorA, scalefactorB
//
template<
class Gemm,
int SchedulerPipelineStageCount_,
class ElementA_,
class ElementB_
>
struct HostCollectiveMainloop<cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpongBlockScaledSm120<SchedulerPipelineStageCount_>,
Gemm, ElementA_, ElementB_> : public
HostCollectiveMainloop<cutlass::gemm::KernelPtrArrayTmaWarpSpecializedBlockScaledSm100<0,0>,
Gemm, ElementA_, ElementB_> {
using Base = HostCollectiveMainloop<cutlass::gemm::KernelPtrArrayTmaWarpSpecializedBlockScaledSm100<0,0>,
Gemm, ElementA_, ElementB_>;
HostCollectiveMainloop(
CheckEquality check_relative_equality_ = CheckEquality::EXACT,
cutlass::Distribution::Kind init_A_ = cutlass::Distribution::Uniform,
cutlass::Distribution::Kind init_B_ = cutlass::Distribution::Uniform,
uint64_t seed_ = Base::kDefaultSeed,
typename Base::LayoutTagA::Stride stride_factor_A_ = typename Base::LayoutTagA::Stride(),
typename Base::LayoutTagB::Stride stride_factor_B_ = typename Base::LayoutTagB::Stride()
) : Base::HostCollectiveMainloop(check_relative_equality_, init_A_, init_B_, seed_, stride_factor_A_, stride_factor_B_) {}
};
//
// Block Scaled Gemm Input Operands : A , B, scalefactorA, scalefactorB
//
template<
class Gemm,
int SchedulerPipelineStageCount_,
class ElementA_,
class ElementB_
>
struct HostCollectiveMainloop<cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeBlockScaledSm120<SchedulerPipelineStageCount_>,
Gemm, ElementA_, ElementB_> : public
HostCollectiveMainloop<cutlass::gemm::KernelPtrArrayTmaWarpSpecializedBlockScaledSm100<0,0>,
Gemm, ElementA_, ElementB_> {
using Base = HostCollectiveMainloop<cutlass::gemm::KernelPtrArrayTmaWarpSpecializedBlockScaledSm100<0,0>,
Gemm, ElementA_, ElementB_>;
HostCollectiveMainloop(
CheckEquality check_relative_equality_ = CheckEquality::EXACT,
cutlass::Distribution::Kind init_A_ = cutlass::Distribution::Uniform,
cutlass::Distribution::Kind init_B_ = cutlass::Distribution::Uniform,
uint64_t seed_ = Base::kDefaultSeed,
typename Base::LayoutTagA::Stride stride_factor_A_ = typename Base::LayoutTagA::Stride(),
typename Base::LayoutTagB::Stride stride_factor_B_ = typename Base::LayoutTagB::Stride()
) : Base::HostCollectiveMainloop(check_relative_equality_, init_A_, init_B_, seed_, stride_factor_A_, stride_factor_B_) {}
};
template<class Gemm>
struct HostCollectiveDefaultEpilogue {
@@ -1529,6 +1579,12 @@ struct HostCollectiveEpilogue {
<< "\n\nComputed Aux =\n" << tensors_Aux[batch].host_view();
}
if constexpr (IsBlockScaleSupported) {
file
<< "\n\nReference SFD =\n" << references_SFD[batch].host_view()
<< "\n\nComputed SFD =\n" << tensors_SFD[batch].host_view();
}
file
<< "\nC =\n" << tensors_C[batch].host_view()
<< "\n\nReference =\n" << references_D[batch].host_view()

6
test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/CMakeLists.txt
View File

@@ -37,6 +37,7 @@ add_custom_target(
cutlass_test_unit_bs_gemm_device_tensorop_epilogue_fusion_sm120
cutlass_test_unit_bs_gemm_device_tensorop_sm120
cutlass_test_unit_bs_gemm_device_tensorop_sm120_stream_k
cutlass_test_unit_bs_grouped_gemm_device_tensorop_sm120
)
cutlass_test_unit_gemm_device_add_executable(
@@ -67,4 +68,9 @@ cutlass_test_unit_gemm_device_add_executable(
sm120_bs_gemm_nvf4_nvf4_f32_f32_stream_k.cu
)
cutlass_test_unit_gemm_device_add_executable(
cutlass_test_unit_bs_grouped_gemm_device_tensorop_sm120
sm120_bs_gemm_nvf4_nvf4_f32_nvf4_group_gemm_fusion.cu
)
endif()

358
test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/sm120_bs_gemm_nvf4_nvf4_f32_nvf4_group_gemm_fusion.cu Normal file
View File

@@ -0,0 +1,358 @@
/***************************************************************************************************
* 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 Tests for device-wide grouped GEMM interface
*/
#include <iostream>
#include "cutlass/cutlass.h"
#include "cute/tensor.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cutlass/numeric_types.h"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/default_epilogue.hpp"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/epilogue/thread/activation.h"
#include "../../../common/cutlass_unit_test.h"
#include "../gemm_testbed_3x_ptr_array.hpp"
using namespace cute;
#if defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED)
// Pingpong kernel schedule
TEST(SM120_Device_Gemm_e2m1t_e2m1n_e2m1t_tensorop_f32_epilogue_VS16_group_pingpong, row_sf) {
using ElementInput = float_e2m1_t;
using ElementA = cutlass::nv_float4_t<ElementInput>;
using ElementB = cutlass::nv_float4_t<ElementInput>;
using ElementC = cutlass::half_t;
using ElementD = cutlass::float_e2m1_t;
using ElementCompute = float;
using ElementAccumulator = float;
using ElementSF = cutlass::float_ue4m3_t;
using ElementSFD = ElementSF;
using ElementAccumulator = float;
using GmemLayoutA = cutlass::layout::RowMajor;
using GmemLayoutB = cutlass::layout::ColumnMajor;
using GmemLayoutC = cutlass::layout::RowMajor;
constexpr int SFVectorSize = 16;
using TileShape_MNK = Shape<_128,_128,_128>;
using ClusterShape_MNK = Shape<_1,_1,_1>;
constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementInput>::value;
constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementInput>::value;
constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
//
// Construct CollectiveEpilogue
//
constexpr int OutputSFVectorSize = SFVectorSize;
// D = alpha * acc + beta * C
// With Row-major BlockScaleFactor generation.
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
OutputSFVectorSize,
ElementD,
ElementCompute,
ElementSFD, GmemLayoutC,
ElementC>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
TileShape_MNK, ClusterShape_MNK,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, GmemLayoutC *, AlignmentC,
ElementD, GmemLayoutC *, AlignmentD,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
//
// Construct CollectiveMainloop
//
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementA, GmemLayoutA *, AlignmentA,
ElementB, GmemLayoutB *, AlignmentB,
ElementAccumulator,
TileShape_MNK, ClusterShape_MNK,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong
>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
cutlass::gemm::GroupProblemShape<Shape<int,int,int>>,
CollectiveMainloop,
CollectiveEpilogue
>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
auto pass = test::gemm::device::TestSmallFusion<Gemm>(1.0, 0.5);
EXPECT_TRUE(pass);
}
TEST(SM120_Device_Gemm_e2m1t_e2m1n_e2m1t_tensorop_f32_epilogue_VS16_group_pingpong, silu_row_sf) {
using ElementInput = float_e2m1_t;
using ElementA = cutlass::nv_float4_t<ElementInput>;
using ElementB = cutlass::nv_float4_t<ElementInput>;
using ElementC = cutlass::half_t;
using ElementD = cutlass::float_e2m1_t;
using ElementCompute = float;
using ElementAccumulator = float;
using ElementSF = cutlass::float_ue4m3_t;
using ElementSFD = ElementSF;
using ElementAccumulator = float;
using GmemLayoutA = cutlass::layout::RowMajor;
using GmemLayoutB = cutlass::layout::ColumnMajor;
using GmemLayoutC = cutlass::layout::RowMajor;
constexpr int SFVectorSize = 16;
using TileShape_MNK = Shape<_128,_128,_256>;
using ClusterShape_MNK = Shape<_1,_1,_1>;
constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementInput>::value;
constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementInput>::value;
constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
//
// Construct CollectiveEpilogue
//
constexpr int OutputSFVectorSize = SFVectorSize;
// D = SiLu(alpha * acc + beta * C)
// With Row-major BlockScaleFactor generation.
using FusionOperation = cutlass::epilogue::fusion::LinCombEltActBlockScaleFactor<
cutlass::epilogue::thread::SiLu,
OutputSFVectorSize,
ElementD,
ElementCompute,
ElementSFD, GmemLayoutC,
ElementC>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
TileShape_MNK, ClusterShape_MNK,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, GmemLayoutC *, AlignmentC,
ElementD, GmemLayoutC *, AlignmentD,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
//
// Construct CollectiveMainloop
//
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementA, GmemLayoutA *, AlignmentA,
ElementB, GmemLayoutB *, AlignmentB,
ElementAccumulator,
TileShape_MNK, ClusterShape_MNK,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong
>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
cutlass::gemm::GroupProblemShape<Shape<int,int,int>>,
CollectiveMainloop,
CollectiveEpilogue
>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
auto pass = test::gemm::device::TestSmallFusion<Gemm>(1.0, 0.5);
EXPECT_TRUE(pass);
}
// Cooperative kenel schedule
TEST(SM120_Device_Gemm_e2m1t_e2m1n_e2m1t_tensorop_f32_epilogue_VS16_group_cooperative, row_sf) {
using ElementInput = float_e2m1_t;
using ElementA = cutlass::nv_float4_t<ElementInput>;
using ElementB = cutlass::nv_float4_t<ElementInput>;
using ElementC = cutlass::half_t;
using ElementD = cutlass::float_e2m1_t;
using ElementCompute = float;
using ElementAccumulator = float;
using ElementSF = cutlass::float_ue4m3_t;
using ElementSFD = ElementSF;
using ElementAccumulator = float;
using GmemLayoutA = cutlass::layout::RowMajor;
using GmemLayoutB = cutlass::layout::ColumnMajor;
using GmemLayoutC = cutlass::layout::RowMajor;
constexpr int SFVectorSize = 16;
using TileShape_MNK = Shape<_128,_128,_128>;
using ClusterShape_MNK = Shape<_1,_1,_1>;
constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementInput>::value;
constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementInput>::value;
constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
//
// Construct CollectiveEpilogue
//
constexpr int OutputSFVectorSize = SFVectorSize;
// D = alpha * acc + beta * C
// With Row-major BlockScaleFactor generation.
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
OutputSFVectorSize,
ElementD,
ElementCompute,
ElementSFD, GmemLayoutC,
ElementC>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
TileShape_MNK, ClusterShape_MNK,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, GmemLayoutC *, AlignmentC,
ElementD, GmemLayoutC *, AlignmentD,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
//
// Construct CollectiveMainloop
//
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementA, GmemLayoutA *, AlignmentA,
ElementB, GmemLayoutB *, AlignmentB,
ElementAccumulator,
TileShape_MNK, ClusterShape_MNK,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::collective::KernelScheduleAuto
>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
cutlass::gemm::GroupProblemShape<Shape<int,int,int>>,
CollectiveMainloop,
CollectiveEpilogue
>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
auto pass = test::gemm::device::TestSmallFusion<Gemm>(1.0, 0.5);
EXPECT_TRUE(pass);
}
TEST(SM120_Device_Gemm_e2m1t_e2m1n_e2m1t_tensorop_f32_epilogue_VS16_group_cooperative, silu_row_sf) {
using ElementInput = float_e2m1_t;
using ElementA = cutlass::nv_float4_t<ElementInput>;
using ElementB = cutlass::nv_float4_t<ElementInput>;
using ElementC = cutlass::half_t;
using ElementD = cutlass::float_e2m1_t;
using ElementCompute = float;
using ElementAccumulator = float;
using ElementSF = cutlass::float_ue4m3_t;
using ElementSFD = ElementSF;
using ElementAccumulator = float;
using GmemLayoutA = cutlass::layout::RowMajor;
using GmemLayoutB = cutlass::layout::ColumnMajor;
using GmemLayoutC = cutlass::layout::RowMajor;
constexpr int SFVectorSize = 16;
using TileShape_MNK = Shape<_128,_128,_256>;
using ClusterShape_MNK = Shape<_1,_1,_1>;
constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementInput>::value;
constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementInput>::value;
constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
//
// Construct CollectiveEpilogue
//
constexpr int OutputSFVectorSize = SFVectorSize;
// D = SiLu(alpha * acc + beta * C)
// With Row-major BlockScaleFactor generation.
using FusionOperation = cutlass::epilogue::fusion::LinCombEltActBlockScaleFactor<
cutlass::epilogue::thread::SiLu,
OutputSFVectorSize,
ElementD,
ElementCompute,
ElementSFD, GmemLayoutC,
ElementC>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
TileShape_MNK, ClusterShape_MNK,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, GmemLayoutC *, AlignmentC,
ElementD, GmemLayoutC *, AlignmentD,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
//
// Construct CollectiveMainloop
//
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementA, GmemLayoutA *, AlignmentA,
ElementB, GmemLayoutB *, AlignmentB,
ElementAccumulator,
TileShape_MNK, ClusterShape_MNK,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::collective::KernelScheduleAuto
>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
cutlass::gemm::GroupProblemShape<Shape<int,int,int>>,
CollectiveMainloop,
CollectiveEpilogue
>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
auto pass = test::gemm::device::TestSmallFusion<Gemm>(1.0, 0.5);
EXPECT_TRUE(pass);
}
#endif // #if defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED)