v3.9 update (#2203)

* v3.9 update

* voidD

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
This commit is contained in:
Yujia Zhai
2025-04-02 12:11:18 -07:00
committed by GitHub
parent 62750a2b75
commit 6f4921858b
129 changed files with 7719 additions and 2036 deletions

View File

@@ -140,7 +140,7 @@ function(cutlass_test_unit_add_executable_split_file NAME)
if (CUTLASS_UNIT_TEST_SPLIT_FILES)
execute_process(
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}
COMMAND ${Python3_EXECUTABLE} ${CUTLASS_SOURCE_DIR}/tools/scripts/split_test_cmake.py
COMMAND ${Python3_EXECUTABLE} ${CUTLASS_SOURCE_DIR}/tools/util/scripts/split_test_cmake.py
${NAME}
${CMAKE_CURRENT_SOURCE_DIR}
--src_files ${SUBARGV}

View File

@@ -499,3 +499,40 @@ TEST(SM80_CuTe_Ampere, CooperativeGemm2_Double_MMA_Predicated_Reg) {
test_cooperative_gemm_col_major_layout_rmem_c<thread_block_size, value_type>(shape_mnk, tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemmLDSMx2) {
constexpr uint32_t thread_block_size = 128;
constexpr int MaxVecBits = 128;
using TA = cute::half_t;
using TB = cute::half_t;
using TC = float;
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
Layout<Shape<_2, _2, _1>, Stride<_1, _2, _0>>,
Tile<_32, _16, _16>
>{};
auto global_a_layout = make_layout(Shape<_32, _32>{}, LayoutRight{});
auto global_b_layout = make_layout(Shape<_16, _32>{}, LayoutRight{});
auto global_c_layout = make_layout(Shape<_32, _16>{}, LayoutRight{});
test_cooperative_gemm<thread_block_size,
MaxVecBits,
TA, TB, TC>
(global_a_layout,
global_b_layout,
global_c_layout,
global_a_layout,
global_b_layout,
global_c_layout,
tiled_mma,
identity{},
identity{},
identity{},
identity{},
SM75_U32x4_LDSM_N{},
SM75_U32x2_LDSM_N{});
}

View File

@@ -188,7 +188,8 @@ template<uint32_t ThreadBlockSize,
class CStoreTransform,
class SMemCopyOpA,
class SMemCopyOpB,
class SMemCopyOpC>
class SMemCopyLdOpC,
class SMemCopyStOpC>
__launch_bounds__(ThreadBlockSize) __global__ void
cooperative_gemm_kernel(GMemALayout gmem_a_layout,
GMemBLayout gmem_b_layout,
@@ -209,7 +210,8 @@ cooperative_gemm_kernel(GMemALayout gmem_a_layout,
CStoreTransform c_store_transform,
SMemCopyOpA a_copy_op,
SMemCopyOpB b_copy_op,
SMemCopyOpC c_copy_op)
SMemCopyLdOpC c_copy_ld_op,
SMemCopyStOpC c_copy_st_op)
{
using namespace cute;
@@ -242,7 +244,7 @@ cooperative_gemm_kernel(GMemALayout gmem_a_layout,
threadIdx.x, tiled_mma,
alpha, s_a_tensor, s_b_tensor, beta, s_c_tensor,
a_load_transform, b_load_transform, c_load_transform, c_store_transform,
a_copy_op, b_copy_op, c_copy_op
a_copy_op, b_copy_op, c_copy_ld_op, c_copy_st_op
);
__syncthreads();
@@ -366,7 +368,8 @@ template<uint32_t ThreadBlockSize,
class CStoreTransform = cute::identity,
class ASMemCopyOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>,
class BSMemCopyOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>,
class CSMemCopyOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>>
class CSMemCopyLdOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>,
class CSMemCopyStOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>>
void test_cooperative_gemm(GMemALayout gmem_a_layout,
GMemBLayout gmem_b_layout,
GMemCLayout gmem_c_layout,
@@ -380,7 +383,8 @@ void test_cooperative_gemm(GMemALayout gmem_a_layout,
CStoreTransform c_store_transform = {},
ASMemCopyOp a_smem_copy_op = {},
BSMemCopyOp b_smem_copy_op = {},
CSMemCopyOp c_smem_copy_op = {})
CSMemCopyLdOp c_smem_copy_ld_op = {},
CSMemCopyStOp c_smem_copy_st_op = {})
{
static_assert(std::is_same_v<typename fp64_tester<TA>::value_type, typename fp64_tester<TB>::value_type>);
static_assert(std::is_same_v<typename fp64_tester<TB>::value_type, typename fp64_tester<TC>::value_type>);
@@ -428,7 +432,7 @@ void test_cooperative_gemm(GMemALayout gmem_a_layout,
TA, TB, TC, decltype(alpha), decltype(beta),
TiledMma,
ALoadTransform, BLoadTransform, CLoadTransform, CStoreTransform,
ASMemCopyOp, BSMemCopyOp, CSMemCopyOp
ASMemCopyOp, BSMemCopyOp, CSMemCopyLdOp, CSMemCopyStOp
>;
ASSERT_EQ(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(shared_memory_size)), 0);
@@ -453,7 +457,8 @@ void test_cooperative_gemm(GMemALayout gmem_a_layout,
c_store_transform,
a_smem_copy_op,
b_smem_copy_op,
c_smem_copy_op
c_smem_copy_ld_op,
c_smem_copy_st_op
);
cudaError_t result = cudaDeviceSynchronize();

View File

@@ -115,3 +115,46 @@ TEST(SM90_CuTe_Hopper, CooperativeGemmTilingF16) {
}
#endif
#if defined(CUTE_ARCH_STSM_SM90_ENABLED)
TEST(SM90_CuTe_Hopper, CooperativeGemmSTSM) {
constexpr uint32_t thread_block_size = 128;
constexpr int MaxVecBits = 128;
using TA = cute::half_t;
using TB = cute::half_t;
using TC = cute::half_t;
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x16_F16F16F16F16_TN>,
Layout<Shape<_2, _2, _1>, Stride<_1, _2, _0>>,
Tile<_32, _32, _16>
>{};
auto global_a_layout = make_layout(Shape<_64, _64>{}, LayoutRight{});
auto global_b_layout = make_layout(Shape<_64, _64>{}, LayoutRight{});
auto global_c_layout = make_layout(Shape<_64, _64>{}, LayoutRight{});
test_cooperative_gemm<thread_block_size,
MaxVecBits,
TA, TB, TC>
(global_a_layout,
global_b_layout,
global_c_layout,
global_a_layout,
global_b_layout,
global_c_layout,
tiled_mma,
identity{},
identity{},
identity{},
identity{},
SM75_U32x4_LDSM_N{},
SM75_U32x4_LDSM_N{},
SM75_U32x4_LDSM_N{},
SM90_U32x4_STSM_N{});
}
#endif

View File

@@ -2449,17 +2449,20 @@ struct HostCollectiveEpilogue {
// example of how to set kernel activation arguments
// see ActivationFunctor::Arguments in activation.h for definition
// if Arguments doesn't exist then fusion_args.activation is empty
auto init_activation_args = [] (auto activation, auto& args) {
using Activation = cute::remove_cvref_t<decltype(activation)>;
if constexpr (cute::is_same_v<Activation, cutlass::epilogue::thread::Clamp<ElementCompute>>) {
args.lower_bound = 0; // Treat Clamp as ReLU
args.upper_bound = cutlass::platform::identity_for_minimum<ElementCompute>();
}
if constexpr (cute::is_same_v<Activation, cutlass::epilogue::thread::ScaledGELU_taylor<ElementCompute>>) {
args.scale = ElementCompute(1);
}
};
if constexpr (cute::is_same_v<ActivationFunctor, cutlass::epilogue::thread::ScaledGELU_taylor<ElementCompute>>) {
fusion_args.activation.scale = ElementCompute(1);
if constexpr (not cute::is_same_v<ActivationFunctor, cutlass::epilogue::thread::Identity<ElementCompute>>) {
init_activation_args(ActivationFunctor{}, fusion_args.activation);
}
// Treat Clamp as ReLU
if constexpr (cute::is_same_v<ActivationFunctor, cutlass::epilogue::thread::Clamp<ElementCompute>>) {
fusion_args.activation.lower_bound = 0;
fusion_args.activation.upper_bound = std::numeric_limits<ElementCompute>::max();
}
if constexpr (IsAbsMaxEnabledD) {
fusion_args.amax_D_ptr = abs_max_D.device_data();
}

View File

@@ -45,24 +45,26 @@ cutlass_test_unit_gemm_device_add_executable(
BATCH_SOURCES ON
BATCH_SIZE 1
sm120_bs_gemm_f4_f4_f32_f32_epilogue_fusion.cu
sm120_bs_gemm_f4_f4_f32_f4_epilogue_fusion.cu
sm120_bs_gemm_f4_f4_f32_bf16_epilogue_fusion.cu
sm120_bs_gemm_nvf4_nvf4_f32_f32_epilogue_fusion.cu
sm120_bs_gemm_nvf4_nvf4_f32_nvf4_epilogue_fusion.cu
sm120_bs_gemm_nvf4_nvf4_f32_bf16_epilogue_fusion.cu
)
cutlass_test_unit_gemm_device_add_executable(
cutlass_test_unit_bs_gemm_device_tensorop_sm120
sm120_bs_gemm_f4_f4_f32_bf16.cu
sm120_bs_gemm_f4_f4_f32_f16.cu
sm120_bs_gemm_f4_f4_f32_f32.cu
sm120_bs_gemm_f4_f4_f32_f32_narrow_output.cu
sm120_bs_gemm_f4_f4_f32_epilogue.cu
sm120_bs_gemm_nvf4_nvf4_f32_bf16.cu
sm120_bs_gemm_nvf4_nvf4_f32_f16.cu
sm120_bs_gemm_nvf4_nvf4_f32_f32.cu
sm120_bs_gemm_nvf4_nvf4_f32_f32_narrow_output.cu
sm120_bs_gemm_nvf4_nvf4_f32_epilogue.cu
sm120_bs_gemm_mxf4_mxf4_f32_f32.cu
sm120_bs_gemm_mxf6_mxf8_f32_f32.cu
)
cutlass_test_unit_gemm_device_add_executable(
cutlass_test_unit_bs_gemm_device_tensorop_sm120_stream_k
sm120_bs_gemm_f4_f4_f32_f32_stream_k.cu
sm120_bs_gemm_nvf4_nvf4_f32_f32_stream_k.cu
)
endif()

View File

@@ -1,590 +0,0 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#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/gemm/collective/collective_builder.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 "../../../common/cutlass_unit_test.h"
#include "../gemm_testbed_3x.hpp"
#if (defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED))
using namespace cute;
///////////////////////////////////////////////////////////////////////////////
namespace kernel_1 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = float;
using ElementD = cutlass::float_e2m1_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue4m3_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::RowMajor;
using LayoutSFD = cutlass::layout::RowMajor;
using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 16 bytes.
static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 16 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_256>;
using ClusterShape = Shape<_1,_1,_1>;
constexpr int SFVectorSize = 16;
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
SFVectorSize,
ElementD,
ElementCompute,
ElementSF,
LayoutSFD
>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
FusionOperation
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
>::CollectiveOp;
template <typename T>
struct dummy {
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
cutlass::gemm::PersistentScheduler>; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_1
namespace kernel_2 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = float;
using ElementD = cutlass::float_e2m1_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue8m0_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::RowMajor;
using LayoutSFD = cutlass::layout::RowMajor;
using ElementPairA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 16 bytes.
static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 16 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_256>;
using ClusterShape = Shape<_1,_1,_1>;
constexpr int SFVectorSize = 32;
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
SFVectorSize,
ElementD,
ElementCompute,
ElementSF,
LayoutSFD
>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
FusionOperation
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedMxf4Sm120
>::CollectiveOp;
template <typename T>
struct dummy {
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
cutlass::gemm::PersistentScheduler>; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_2
namespace kernel_3 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = cutlass::half_t;
using ElementD = cutlass::float_e2m1_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue4m3_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::RowMajor;
using LayoutSFD = cutlass::layout::RowMajor;
using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 16 bytes.
static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 16 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_256>;
using ClusterShape = Shape<_1,_1,_1>;
constexpr int SFVectorSize = 16;
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
SFVectorSize,
ElementD,
ElementCompute,
ElementSF,
LayoutSFD
>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
FusionOperation
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
>::CollectiveOp;
template <typename T>
struct dummy {
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
cutlass::gemm::PersistentScheduler>; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_3
namespace kernel_4 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = cutlass::half_t;
using ElementD = cutlass::float_e2m1_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue8m0_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::RowMajor;
using LayoutSFD = cutlass::layout::RowMajor;
using ElementPairA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 16 bytes.
static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 16 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_256>;
using ClusterShape = Shape<_1,_1,_1>;
constexpr int SFVectorSize = 32;
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
SFVectorSize,
ElementD,
ElementCompute,
ElementSF,
LayoutSFD
>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
FusionOperation
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedMxf4Sm120
>::CollectiveOp;
template <typename T>
struct dummy {
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
cutlass::gemm::PersistentScheduler>; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_4
namespace kernel_5 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = cutlass::bfloat16_t;
using ElementD = cutlass::float_e2m1_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue4m3_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::RowMajor;
using LayoutSFD = cutlass::layout::RowMajor;
using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 16 bytes.
static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 16 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_256>;
using ClusterShape = Shape<_1,_1,_1>;
constexpr int SFVectorSize = 16;
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
SFVectorSize,
ElementD,
ElementCompute,
ElementSF,
LayoutSFD
>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
FusionOperation
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
>::CollectiveOp;
template <typename T>
struct dummy {
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
cutlass::gemm::PersistentScheduler>; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_5
namespace kernel_6 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = cutlass::bfloat16_t;
using ElementD = cutlass::float_e2m1_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue8m0_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::RowMajor;
using LayoutSFD = cutlass::layout::RowMajor;
using ElementPairA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 16 bytes.
static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 16 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_256>;
using ClusterShape = Shape<_1,_1,_1>;
constexpr int SFVectorSize = 32;
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
SFVectorSize,
ElementD,
ElementCompute,
ElementSF,
LayoutSFD
>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
FusionOperation
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedMxf4Sm120
>::CollectiveOp;
template <typename T>
struct dummy {
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
cutlass::gemm::PersistentScheduler>; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_6
namespace kernel_7 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = void;
using ElementD = cutlass::float_e2m1_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue4m3_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::RowMajor;
using LayoutSFD = cutlass::layout::RowMajor;
using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 16 bytes.
static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 16 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementD>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_256>;
using ClusterShape = Shape<_1,_1,_1>;
constexpr int SFVectorSize = 16;
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
SFVectorSize,
ElementD,
ElementCompute,
ElementSF,
LayoutSFD,
ElementC
>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
FusionOperation
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
>::CollectiveOp;
template <typename T>
struct dummy {
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
cutlass::gemm::PersistentScheduler>; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_7
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs16_tensor_op_f32_f32_epilogue_vs16, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_1::Gemm>(1.0, 0.5);
EXPECT_TRUE(result);
}
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs32_tensor_op_f32_f32_epilogue_vs32, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_2::Gemm>(1.0, 0.5);
EXPECT_TRUE(result);
}
// ==== mixed datatypes for C (fp16/bf16) / D (fp32) matrices ==== //
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs16_tensor_op_f16_f32_epilogue_vs16, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_3::Gemm>(1.0, 0.5);
EXPECT_TRUE(result);
}
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs32_tensor_op_f16_f32_epilogue_vs32, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_4::Gemm>(1.0, 0.5);
EXPECT_TRUE(result);
}
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs16_tensor_op_bf16_f32_epilogue_vs16, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_5::Gemm>(1.0, 0.5);
EXPECT_TRUE(result);
}
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs32_tensor_op_bf16_f32_epilogue_vs32, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_6::Gemm>(1.0, 0.5);
EXPECT_TRUE(result);
}
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs32_tensor_op_void_f32_epilogue_vs32, 128x128x256) {
bool result = test::gemm::device::TestSmallFusion<kernel_7::Gemm>(1.0, 0.0);
EXPECT_TRUE(result);
}
#endif // (defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED))

View File

@@ -87,7 +87,7 @@ namespace kernel_1 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -97,7 +97,7 @@ namespace kernel_1 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedMxf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -114,7 +114,8 @@ namespace kernel_1 {
} // kernel_1
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs32_tensor_op_f32_static_sched, 128x128x256) {
TEST(SM120_Device_Blockscaled_Gemm_mxf4t_mxf4n_f32n_tensor_op_f32, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_1::Gemm, true>(1.0, 0.5);
EXPECT_TRUE(result);
}

View File

@@ -0,0 +1,123 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#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 "../../../common/cutlass_unit_test.h"
#include "../gemm_testbed_3x.hpp"
#if (defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED))
using namespace cute;
///////////////////////////////////////////////////////////////////////////////
namespace kernel_1 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = float;
using ElementD = float;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue8m0_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::ColumnMajor;
using ElementPairA = cutlass::mx_float6_t<cutlass::float_e3m2_t>;
using ElementPairB = cutlass::mx_float8_t<cutlass::float_e5m2_t>;
static constexpr int AlignmentA = 64 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 64 bytes.
static constexpr int AlignmentB = 96 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 96 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_128>;
using ClusterShape = Shape<_1,_1,_1>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
struct dummy {
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue
>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_1
TEST(SM120_Device_Blockscaled_Gemm_mxf6t_mxf8n_f32n_tensor_op_f32, 128x128x128) {
bool result = test::gemm::device::TestSmall<kernel_1::Gemm, true>(1.0, 0.5);
EXPECT_TRUE(result);
}
#endif // (defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED))

View File

@@ -87,7 +87,7 @@ namespace kernel_1 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -97,7 +97,7 @@ namespace kernel_1 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -115,74 +115,10 @@ namespace kernel_1 {
} // kernel_1
namespace kernel_3 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = cutlass::bfloat16_t;
using ElementD = cutlass::bfloat16_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue8m0_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::ColumnMajor;
using ElementPairA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 64 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 64 bytes.
static constexpr int AlignmentB = 64 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 64 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_128>;
using ClusterShape = Shape<_1,_1,_1>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedMxf8f6f4Sm120
>::CollectiveOp;
template <typename T>
struct dummy {
using TileSchedulerTag = cutlass::gemm::PersistentScheduler; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
TileSchedulerTag>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_3
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_bf16, 128x128x256) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_f32n_tensor_op_bf16, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_1::Gemm, true>(1.0, 0.5);
EXPECT_TRUE(result);
}
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs32_tensor_op_bf16, 128x128x128) {
bool result = test::gemm::device::TestSmall<kernel_3::Gemm, true>(1.0, 0.5);
EXPECT_TRUE(result);
}
#endif // (defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED))

View File

@@ -97,7 +97,7 @@ namespace kernel_1 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -108,7 +108,7 @@ namespace kernel_1 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -167,7 +167,7 @@ namespace kernel_2 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -178,7 +178,7 @@ namespace kernel_2 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -237,7 +237,7 @@ namespace kernel_3 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -248,7 +248,7 @@ namespace kernel_3 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -309,7 +309,7 @@ namespace kernel_4 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -320,7 +320,7 @@ namespace kernel_4 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -343,7 +343,7 @@ namespace kernel_4 {
// Acc: fp32
// Scale (alpha, beta): fp32
// D: bf16
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_bf16n_vs16_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias_relu) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_bf16n_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias_relu) {
bool result = test::gemm::device::TestSmallFusion<kernel_1::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}
@@ -354,7 +354,7 @@ TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_bf16n_vs16_tensor_op_f32_f32_ep
// Acc: fp32
// Scale (alpha, beta): fp32
// D: bf16
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_bf16n_vs16_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias_gelu) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_bf16n_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias_gelu) {
bool result = test::gemm::device::TestSmallFusion<kernel_2::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}
@@ -365,7 +365,7 @@ TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_bf16n_vs16_tensor_op_f32_f32_ep
// Acc: fp32
// Scale (alpha, beta): fp32
// D: bf16
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_bf16n_vs16_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_row_bias_relu) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_bf16n_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_row_bias_relu) {
bool result = test::gemm::device::TestSmallFusion<kernel_3::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}
@@ -377,7 +377,7 @@ TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_bf16n_vs16_tensor_op_f32_f32_ep
// Acc: fp32
// Scale (alpha, beta): fp32
// D: bf16
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_bf16n_vs16_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_row_bias_gelu) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_bf16n_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_row_bias_gelu) {
bool result = test::gemm::device::TestSmallFusion<kernel_4::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}

View File

@@ -39,6 +39,7 @@
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/default_epilogue.hpp"
@@ -58,8 +59,8 @@ using namespace cute;
namespace kernel_1 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = cutlass::half_t;
using ElementD = cutlass::half_t;
using ElementC = float;
using ElementD = cutlass::float_e2m1_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue4m3_t;
@@ -67,7 +68,8 @@ namespace kernel_1 {
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::RowMajor;
using LayoutSFD = cutlass::layout::RowMajor;
using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
@@ -80,6 +82,15 @@ namespace kernel_1 {
using TileShape = Shape<_128,_128,_256>;
using ClusterShape = Shape<_1,_1,_1>;
constexpr int SFVectorSize = 16;
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
SFVectorSize,
ElementD,
ElementCompute,
ElementSF,
LayoutSFD
>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
@@ -87,7 +98,8 @@ namespace kernel_1 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -97,40 +109,41 @@ namespace kernel_1 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
struct dummy {
using TileSchedulerTag = cutlass::gemm::PersistentScheduler; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
TileSchedulerTag>;
cutlass::gemm::PersistentScheduler>; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_1
namespace kernel_2 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = cutlass::half_t;
using ElementD = cutlass::half_t;
using ElementD = cutlass::float_e2m1_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue8m0_t;
using ElementSF = cutlass::float_ue4m3_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::RowMajor;
using LayoutSFD = cutlass::layout::RowMajor;
using ElementPairA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 16 bytes.
static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 16 bytes.
@@ -140,6 +153,15 @@ namespace kernel_2 {
using TileShape = Shape<_128,_128,_256>;
using ClusterShape = Shape<_1,_1,_1>;
constexpr int SFVectorSize = 16;
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
SFVectorSize,
ElementD,
ElementCompute,
ElementSF,
LayoutSFD
>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
@@ -147,7 +169,8 @@ namespace kernel_2 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -157,97 +180,189 @@ namespace kernel_2 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedMxf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
struct dummy {
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue
>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_2
namespace kernel_3 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = cutlass::half_t;
using ElementD = cutlass::half_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue8m0_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::ColumnMajor;
using ElementPairA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 64 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 64 bytes.
static constexpr int AlignmentB = 64 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 64 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_128>;
using ClusterShape = Shape<_1,_1,_1>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedMxf8f6f4Sm120
>::CollectiveOp;
template <typename T>
struct dummy {
using TileSchedulerTag = cutlass::gemm::PersistentScheduler; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
TileSchedulerTag>;
cutlass::gemm::PersistentScheduler>; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_2
namespace kernel_3 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = cutlass::bfloat16_t;
using ElementD = cutlass::float_e2m1_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue4m3_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::RowMajor;
using LayoutSFD = cutlass::layout::RowMajor;
using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 16 bytes.
static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 16 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_256>;
using ClusterShape = Shape<_1,_1,_1>;
constexpr int SFVectorSize = 16;
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
SFVectorSize,
ElementD,
ElementCompute,
ElementSF,
LayoutSFD
>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
struct dummy {
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
cutlass::gemm::PersistentScheduler>; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_3
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_f16, 128x128x256) {
namespace kernel_4 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = void;
using ElementD = cutlass::float_e2m1_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue4m3_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::RowMajor;
using LayoutSFD = cutlass::layout::RowMajor;
using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 16 bytes.
static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 16 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementD>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_256>;
using ClusterShape = Shape<_1,_1,_1>;
constexpr int SFVectorSize = 16;
using FusionOperation = cutlass::epilogue::fusion::LinCombBlockScaleFactor<
SFVectorSize,
ElementD,
ElementCompute,
ElementSF,
LayoutSFD,
ElementC
>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
struct dummy {
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
cutlass::gemm::PersistentScheduler>; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_4
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_nvf4t_tensor_op_f32_f32_epilogue_vs16, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_1::Gemm>(1.0, 0.5);
EXPECT_TRUE(result);
}
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs32_tensor_op_f16_static_sched, 128x128x256) {
// ==== mixed datatypes for C (fp16/bf16) / D (fp32) matrices ==== //
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_nvf4t_tensor_op_f16_f32_epilogue_vs16, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_2::Gemm>(1.0, 0.5);
EXPECT_TRUE(result);
}
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs32_tensor_op_f16, 128x128x128) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_nvf4t_tensor_op_bf16_f32_epilogue_vs16, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_3::Gemm>(1.0, 0.5);
EXPECT_TRUE(result);
}
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_nvf4t_tensor_op_void_f32_epilogue_vs32, 128x128x256) {
bool result = test::gemm::device::TestSmallFusion<kernel_4::Gemm>(1.0, 0.0);
EXPECT_TRUE(result);
}
#endif // (defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED))

View File

@@ -0,0 +1,126 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#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 "../../../common/cutlass_unit_test.h"
#include "../gemm_testbed_3x.hpp"
#if (defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED))
using namespace cute;
///////////////////////////////////////////////////////////////////////////////
namespace kernel_1 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = cutlass::half_t;
using ElementD = cutlass::half_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue4m3_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::ColumnMajor;
using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 16 bytes.
static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 16 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_256>;
using ClusterShape = Shape<_1,_1,_1>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
struct dummy {
using TileSchedulerTag = cutlass::gemm::PersistentScheduler; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
TileSchedulerTag>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_1
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_f32n_tensor_op_f16, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_1::Gemm>(1.0, 0.5);
EXPECT_TRUE(result);
}
#endif // (defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED))

View File

@@ -0,0 +1,125 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#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 "../../../common/cutlass_unit_test.h"
#include "../gemm_testbed_3x.hpp"
#if (defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED))
using namespace cute;
///////////////////////////////////////////////////////////////////////////////
namespace kernel_1 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = float;
using ElementD = float;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue4m3_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::ColumnMajor;
using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 16 bytes.
static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 16 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_256>;
using ClusterShape = Shape<_1,_1,_1>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedPingpong
>::CollectiveOp;
template <typename T>
struct dummy {
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue
>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_1
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_f32n_tensor_op_f32, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_1::Gemm, true>(1.0, 0.5);
EXPECT_TRUE(result);
}
#endif // (defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED))

View File

@@ -97,7 +97,7 @@ namespace kernel_1 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -108,7 +108,7 @@ namespace kernel_1 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -167,7 +167,7 @@ namespace kernel_2 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -178,7 +178,7 @@ namespace kernel_2 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -237,7 +237,7 @@ namespace kernel_3 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -248,7 +248,7 @@ namespace kernel_3 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -307,7 +307,7 @@ namespace kernel_4 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -318,7 +318,7 @@ namespace kernel_4 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -377,7 +377,7 @@ namespace kernel_5 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -388,7 +388,7 @@ namespace kernel_5 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -447,7 +447,7 @@ namespace kernel_6 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -458,7 +458,7 @@ namespace kernel_6 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -481,7 +481,7 @@ namespace kernel_6 {
// Acc: fp32
// Scale (alpha, beta): fp32
// D: fp32
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_f32_f32_epilogue, 128x128x256_per_row_bias_relu) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_f32n_tensor_op_f32_f32_epilogue, 128x128x256_per_row_bias_relu) {
bool result = test::gemm::device::TestSmallFusion<kernel_1::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}
@@ -492,7 +492,7 @@ TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_f32_f32_epi
// Acc: fp32
// Scale (alpha, beta): fp32
// D: fp32
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_f32_f32_epilogue, 128x128x256_per_row_bias_gelu) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_f32n_tensor_op_f32_f32_epilogue, 128x128x256_per_row_bias_gelu) {
bool result = test::gemm::device::TestSmallFusion<kernel_2::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}
@@ -503,7 +503,7 @@ TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_f32_f32_epi
// Acc: fp32
// Scale (alpha, beta): fp32
// D: fp32
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias_gelu) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_f32n_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias_gelu) {
bool result = test::gemm::device::TestSmallFusion<kernel_3::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}
@@ -514,7 +514,7 @@ TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_f32_f32_epi
// Acc: fp32
// Scale (alpha, beta): fp32
// D: fp32
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias_relu) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_f32n_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias_relu) {
bool result = test::gemm::device::TestSmallFusion<kernel_4::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}
@@ -525,7 +525,7 @@ TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_f32_f32_epi
// Acc: fp32
// Scale (alpha, beta): fp32
// D: fp32
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_f32_f32_epilogue, 128x128x256_per_row_bias_clamp) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_f32n_tensor_op_f32_f32_epilogue, 128x128x256_per_row_bias_clamp) {
bool result = test::gemm::device::TestSmallFusion<kernel_5::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}
@@ -536,7 +536,7 @@ TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_f32_f32_epi
// Acc: fp32
// Scale (alpha, beta): fp32
// D: fp32
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias_clamp) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_f32n_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias_clamp) {
bool result = test::gemm::device::TestSmallFusion<kernel_6::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}

View File

@@ -102,7 +102,7 @@ namespace kernel_1 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -113,7 +113,7 @@ namespace kernel_1 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -130,89 +130,9 @@ namespace kernel_1 {
} // kernel_1
namespace kernel_2 {
using ElementA = cutlass::float_e2m1_t;
using ElementB = cutlass::float_e2m1_t;
using ElementC = cutlass::bfloat16_t;
using ElementD = cutlass::float_e2m3_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementSF = cutlass::float_ue8m0_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::RowMajor;
using LayoutD = cutlass::layout::RowMajor;
using ElementPairA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
using ElementPairB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value; // Align to 16 bytes.
static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value; // Align to 16 bytes.
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
using TileShape = Shape<_128,_128,_256>;
using ClusterShape = Shape<_1,_1,_1>;
constexpr int SFVectorSize = 32;
using LayoutSFD = cutlass::layout::RowMajor;
using ElementBias = cutlass::bfloat16_t;
using GmemLayoutSFC = cutlass::layout::RowMajor;
using FusionOperation = cutlass::epilogue::fusion::LinCombPerColBiasEltActBlockScaleFactor<
cutlass::epilogue::thread::ReLU,
SFVectorSize,
ElementD,
ElementCompute,
ElementSF, LayoutSFD,
ElementBias,
ElementC>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
,FusionOperation
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
ElementPairA, LayoutA, AlignmentA,
ElementPairB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedMxf4Sm120
>::CollectiveOp;
template <typename T>
struct dummy {
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
cutlass::gemm::PersistentScheduler>; // both void (default) and PersistentScheduler map to dynamic scheduler with CLC query
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GemmKernel = typename dummy<void>::GemmKernel;
using Gemm = typename dummy<void>::Gemm;
} // kernel_2
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_f32_fe2m3n, 128x128x256) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_f32n_tensor_op_f32_fe2m3n, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_1::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs32_tensor_op_f32_fe2m3n, 128x128x256) {
bool result = test::gemm::device::TestSmallFusion<kernel_2::Gemm, false, false>(1.0, 0);
EXPECT_TRUE(result);
}
#endif // (defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED))

View File

@@ -87,7 +87,7 @@ namespace kernel_1 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -97,7 +97,7 @@ namespace kernel_1 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -115,7 +115,7 @@ namespace kernel_1 {
} // kernel_1
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_f32n_vs16_tensor_op_f32_stream_k, 128x128x256) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_f32n_tensor_op_f32_stream_k, 128x128x256) {
bool result = test::gemm::device::TestSmall<kernel_1::Gemm, true>(1.0, 0.5);
EXPECT_TRUE(result);
}

View File

@@ -104,7 +104,7 @@ namespace kernel_1 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -115,7 +115,7 @@ namespace kernel_1 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -182,7 +182,7 @@ namespace kernel_2 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -193,7 +193,7 @@ namespace kernel_2 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -260,7 +260,7 @@ namespace kernel_3 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -271,7 +271,7 @@ namespace kernel_3 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -336,7 +336,7 @@ namespace kernel_4 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -347,7 +347,7 @@ namespace kernel_4 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -413,7 +413,7 @@ namespace kernel_5 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -424,7 +424,7 @@ namespace kernel_5 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -490,7 +490,7 @@ namespace kernel_6 {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative,
cutlass::epilogue::collective::EpilogueScheduleAuto,
FusionOperation
>::CollectiveOp;
@@ -501,7 +501,7 @@ namespace kernel_6 {
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120
cutlass::gemm::KernelTmaWarpSpecializedCooperative
>::CollectiveOp;
template <typename T>
@@ -527,7 +527,7 @@ namespace kernel_6 {
// Acc: fp32
// Scale (alpha, beta): fp32
// D: bf16
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs16_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_row_bias) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_nvf4t_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_row_bias) {
bool result = test::gemm::device::TestSmallFusion<kernel_1::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}
@@ -538,7 +538,7 @@ TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs16_tensor_op_f32_f32_e
// Acc: fp32
// Scale (alpha, beta): fp32
// D: bf16
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs16_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_row_bias_relu) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_nvf4t_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_row_bias_relu) {
bool result = test::gemm::device::TestSmallFusion<kernel_2::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}
@@ -549,7 +549,7 @@ TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs16_tensor_op_f32_f32_e
// Acc: fp32
// Scale (alpha, beta): fp32
// D: bf16
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs16_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_row_bias_gelu) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_nvf4t_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_row_bias_gelu) {
bool result = test::gemm::device::TestSmallFusion<kernel_3::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}
@@ -560,7 +560,7 @@ TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs16_tensor_op_f32_f32_e
// Acc: fp32
// Scale (alpha, beta): fp32
// D: bf16
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs16_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_nvf4t_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias) {
bool result = test::gemm::device::TestSmallFusion<kernel_4::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}
@@ -571,7 +571,7 @@ TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs16_tensor_op_f32_f32_e
// Acc: fp32
// Scale (alpha, beta): fp32
// D: bf16
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs16_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias_relu) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_nvf4t_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias_relu) {
bool result = test::gemm::device::TestSmallFusion<kernel_5::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}
@@ -582,7 +582,7 @@ TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs16_tensor_op_f32_f32_e
// Acc: fp32
// Scale (alpha, beta): fp32
// D: bf16
TEST(SM120_Device_Blockscaled_Gemm_fe2m1t_fe2m1n_fe2m1t_vs16_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias_gelu) {
TEST(SM120_Device_Blockscaled_Gemm_nvf4t_nvf4n_nvf4t_tensor_op_f32_f32_epilogue, 128x128x256_alpha_beta_per_col_bias_gelu) {
bool result = test::gemm::device::TestSmallFusion<kernel_6::Gemm, false, false>(1.0, 0.5);
EXPECT_TRUE(result);
}

View File

@@ -83,7 +83,7 @@ TEST(SM120_Device_Gemm_fe2m1t_fe2m1n_f16n_void_f32_tensor_op, 128x64x128_1x1x1)
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -134,7 +134,7 @@ TEST(SM120_Device_Gemm_fe2m1t_fe2m1n_f16n_void_f16_tensor_op, 128x64x128_1x1x1)
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -185,7 +185,7 @@ TEST(SM120_Device_Gemm_fe2m1t_fe2m1n_f16n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -236,7 +236,7 @@ TEST(SM120_Device_Gemm_fe2m1t_fe2m1n_f16n_tensor_op_f16, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

View File

@@ -83,7 +83,7 @@ TEST(SM120_Device_Gemm_fe2m1t_fe2m1n_f32n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

View File

@@ -83,7 +83,7 @@ TEST(SM120_Device_Gemm_fe2m1t_fe3m2n_f16n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -134,7 +134,7 @@ TEST(SM120_Device_Gemm_fe2m3t_fe2m1n_f16n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

View File

@@ -85,7 +85,7 @@ TEST(SM120_Device_Gemm_fe2m1t_fe3m2n_f16n_tensor_op_fe2m3n, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -138,7 +138,7 @@ TEST(SM120_Device_Gemm_fe2m3t_fe2m1n_f16n_tensor_op_fe2m1t, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

View File

@@ -83,7 +83,7 @@ TEST(SM120_Device_Gemm_fe2m1t_fe3m2n_f32n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -134,7 +134,7 @@ TEST(SM120_Device_Gemm_fe2m3t_fe2m1n_f32n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

View File

@@ -85,7 +85,7 @@ TEST(SM120_Device_Gemm_fe2m1t_fe3m2n_f32n_tensor_op_fe2m3n, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -137,7 +137,7 @@ TEST(SM120_Device_Gemm_fe2m3t_fe2m1n_f32n_tensor_op_fe2m1t, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

View File

@@ -83,7 +83,7 @@ TEST(SM120_Device_Gemm_fe2m1t_fe5m2n_f16n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -134,7 +134,7 @@ TEST(SM120_Device_Gemm_fe4m3t_fe2m1n_f16n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

View File

@@ -83,7 +83,7 @@ TEST(SM120_Device_Gemm_fe2m1t_fe5m2n_f32n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -134,7 +134,7 @@ TEST(SM120_Device_Gemm_fe4m3t_fe2m1n_f32n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

View File

@@ -83,7 +83,7 @@ TEST(SM120_Device_Gemm_fe3m2t_fe3m2n_f16n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

View File

@@ -83,7 +83,7 @@ TEST(SM120_Device_Gemm_fe3m2t_fe3m2n_f32n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

View File

@@ -83,7 +83,7 @@ TEST(SM120_Device_Gemm_fe3m2t_fe4m3n_f16n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -134,7 +134,7 @@ TEST(SM120_Device_Gemm_fe4m3t_fe3m2n_f16n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

View File

@@ -83,7 +83,7 @@ TEST(SM120_Device_Gemm_fe3m2t_fe4m3n_f32n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -134,7 +134,7 @@ TEST(SM120_Device_Gemm_fe4m3t_fe3m2n_f32n_tensor_op_f32, 128x64x128_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

View File

@@ -82,7 +82,7 @@ TEST(SM120_Device_Gemm_fe4m3t_fe4m3n_f16n_tensor_op_f32, 128x64x64_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

View File

@@ -82,7 +82,7 @@ TEST(SM120_Device_Gemm_fe4m3t_fe4m3n_f32n_tensor_op_f32, 128x64x64_1x1x1) {
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
cutlass::epilogue::TmaWarpSpecializedCooperative
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

View File

@@ -137,7 +137,7 @@ public:
return true;
#endif
#if 1
#if 0
bool is_success = false;
for (int i = 0; i< 10; i++){
printf("iteration = %d\n", i);