v3.9 update (#2203)
* v3.9 update * voidD --------- Co-authored-by: yuzhai <yuzhai@nvidia.com>
This commit is contained in:
@@ -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}
|
||||
|
||||
@@ -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{});
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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))
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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))
|
||||
@@ -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))
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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))
|
||||
@@ -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))
|
||||
@@ -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))
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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))
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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<
|
||||
|
||||
@@ -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<
|
||||
|
||||
@@ -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<
|
||||
|
||||
@@ -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<
|
||||
|
||||
@@ -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<
|
||||
|
||||
@@ -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<
|
||||
|
||||
@@ -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<
|
||||
|
||||
@@ -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<
|
||||
|
||||
@@ -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<
|
||||
|
||||
@@ -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<
|
||||
|
||||
@@ -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<
|
||||
|
||||
@@ -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<
|
||||
|
||||
@@ -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<
|
||||
|
||||
@@ -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<
|
||||
|
||||
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user