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

Small Tile N BlockScaled GEMM + Grouped GEMM (#3176)

Browse Source 
Co-authored-by: dePaul Miller <23461061+depaulmillz@users.noreply.github.com>
This commit is contained in:
dePaul Miller
2026-04-21 12:32:40 -04:00
committed by GitHub
co-authored by dePaul Miller
parent aeba0d3723
commit b46b16d003
13 changed files with 468 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/gemm/collective/builders/sm120_blockscaled_mma_builder.inl
+4 -2
View File
@@ -199,9 +199,11 @@ struct CollectiveBuilder<
using sSFA_stride = decltype(make_stride(sSFA_strideM{}, sSFA_strideK{}));
using SmemLayoutAtomSFA = decltype(make_layout( sSFA_shape{}, sSFA_stride{}));
using sSFB_shapeN = decltype(prepend(size<1>(TileShape_MNK{}) / Blk_MN{}, mnBasicBlockShape{}));
using sSFBTileShape_N = Int<cute::max(size<1>(TileShape_MNK{}), 128)>;
using sSFB_shapeN = decltype(prepend(sSFBTileShape_N{} / Blk_MN{}, mnBasicBlockShape{}));
using sSFB_strideN = sSF_strideMN;
using sSFB_strideK = decltype(prepend(make_stride(Int<MMA_NSF>{}, size<1>(TileShape_MNK{}) / Blk_MN{} * Blk_Elems{}), kBasicBlockStride{}));
using sSFB_strideK = decltype(prepend(make_stride(Int<MMA_NSF>{}, sSFBTileShape_N{} / Blk_MN{} * Blk_Elems{}), kBasicBlockStride{}));
using sSFB_shape = decltype(make_shape( sSFB_shapeN{}, sSF_shapeK{}));
using sSFB_stride = decltype(make_stride(sSFB_strideN{}, sSFB_strideK{}));
using SmemLayoutAtomSFB = decltype(make_layout( sSFB_shape{}, sSFB_stride{}));
include/cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp
+35 -8
View File
@@ -203,11 +203,22 @@ struct CollectiveMma<
append(stride(SmemLayoutAtomSFA{}), size(filter_zeros(SmemLayoutAtomSFA{})))
));
using SmemLayoutSFB = decltype(make_layout(
using SmemLayoutSFB_ = decltype(make_layout(
append(shape(SmemLayoutAtomSFB{}), Int<DispatchPolicy::Stages>{}),
append(stride(SmemLayoutAtomSFB{}), size(filter_zeros(SmemLayoutAtomSFB{})))
));
using TileShapeSFB = cute::conditional_t<size<1>(TileShape{}) < 128,
decltype(cute::make_shape(
shape<0>(TileShape{}),
Int<128>{},
shape<2>(TileShape{}))),
TileShape>;
using SmemLayoutSFB = cute::conditional_t<size<1>(TileShape{}) < 128,
decltype(cute::logical_divide(SmemLayoutSFB_{}, select<1,2>(TileShape{}))),
SmemLayoutSFB_>;
static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3.");
static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3.");
@@ -321,7 +332,7 @@ struct CollectiveMma<
GmemTiledCopySFB{},
make_tensor(static_cast<ElementSF const*>(nullptr), InternalLayoutSFB{}),
SmemLayoutSFB{}(_,_,cute::Int<0>{}),
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
make_shape(shape<1>(TileShapeSFB{}), shape<2>(TileShapeSFB{})),
_1{})); // No programmatic multicast
TMA_A tma_load_a;
@@ -417,7 +428,7 @@ struct CollectiveMma<
GmemTiledCopySFB{},
tensor_sfb,
SmemLayoutSFB{}(_,_,cute::Int<0>{}),
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
make_shape(shape<1>(TileShapeSFB{}), shape<2>(TileShapeSFB{})),
_1{}); // No programmatic multicast
return {
@@ -712,10 +723,15 @@ struct CollectiveMma<
// Partition the inputs based on the current block coordinates.
auto [m_coord, n_coord, k_coord, l_coord] = blk_coord;
auto broadcast_n = make_layout(
make_shape(Int<size<1>(TileShapeSFB{}) / size<1>(TileShape{})>{},
Int<cute::numeric_limits<int>::max()>{}),
make_stride(_0{}, size<1>(TileShapeSFB{}) / size<1>(TileShape{})));
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
Tensor gSFA = gSFA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
Tensor gSFB = gSFB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
Tensor gSFB = gSFB_nkl(_,_,broadcast_n(n_coord),_,l_coord); // (BLK_N,BLK_K,k)
// Partition source and destination tensors for tma copies
Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k)
@@ -778,7 +794,8 @@ struct CollectiveMma<
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class FrgTensorC
class FrgTensorC,
class BlockCoord
>
CUTLASS_DEVICE void
mma(MainloopPipeline pipeline,
@@ -787,7 +804,8 @@ struct CollectiveMma<
int k_tile_count,
int thread_idx,
TensorStorage& shared_tensors,
[[maybe_unused]] Params const& params) {
[[maybe_unused]] Params const& params,
BlockCoord const& blk_coord) {
using namespace cute;
static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
@@ -796,8 +814,17 @@ struct CollectiveMma<
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
Tensor sSFA = make_tensor(make_smem_ptr(shared_tensors.smem_SFA.begin()), SmemLayoutSFA{}); // (BLK_M,BLK_K,PIPE)
Tensor sSFB = make_tensor(make_smem_ptr(shared_tensors.smem_SFB.begin()), SmemLayoutSFB{}); // (BLK_N,BLK_K,PIPE)
Tensor sSFA = make_tensor(make_smem_ptr(shared_tensors.smem_SFA.begin()), SmemLayoutSFA{}); // (BLK_M,BLK_K,PIPE)
Tensor sSFB = [&]() {
if constexpr (size<1>(TileShape{}) >= 128) {
return make_tensor(make_smem_ptr(shared_tensors.smem_SFB.begin()), SmemLayoutSFB{}); // (BLK_N,BLK_K,PIPE)
}
else {
Tensor temp = make_tensor(make_smem_ptr(shared_tensors.smem_SFB.begin()), SmemLayoutSFB{}); // (BLK_SFB_N,BLK_K,PIPE)
auto n = get<1>(blk_coord);
return temp(make_coord(_,n % (size<1>(TileShapeSFB{}) / size<1>(TileShape{}))), _, _);
}
}();
//
// Define C accumulators and A/B partitioning
include/cutlass/gemm/collective/sm120_blockscaled_mma_tma.hpp
+36 -10
View File
@@ -198,11 +198,22 @@ struct CollectiveMma<
append(stride(SmemLayoutAtomSFA{}), size(filter_zeros(SmemLayoutAtomSFA{})))
));
using SmemLayoutSFB = decltype(make_layout(
using SmemLayoutSFB_ = decltype(make_layout(
append(shape(SmemLayoutAtomSFB{}), Int<DispatchPolicy::Stages>{}),
append(stride(SmemLayoutAtomSFB{}), size(filter_zeros(SmemLayoutAtomSFB{})))
));
using TileShapeSFB = cute::conditional_t<size<1>(TileShape{}) < 128,
decltype(cute::make_shape(
shape<0>(TileShape{}),
Int<128>{},
shape<2>(TileShape{}))),
TileShape>;
using SmemLayoutSFB = cute::conditional_t<size<1>(TileShape{}) < 128,
decltype(cute::logical_divide(SmemLayoutSFB_{}, select<1,2>(TileShape{}))),
SmemLayoutSFB_>;
static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3.");
static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3.");
@@ -307,7 +318,7 @@ struct CollectiveMma<
GmemTiledCopySFB{},
make_tensor(static_cast<ElementSF const*>(nullptr), LayoutSFB{}),
SmemLayoutSFB{}(_,_,cute::Int<0>{}),
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
make_shape(shape<1>(TileShapeSFB{}), shape<2>(TileShapeSFB{})),
_1{})); // No programmatic multicast
TMA_A tma_load_a;
@@ -367,7 +378,7 @@ struct CollectiveMma<
GmemTiledCopySFB{},
tensor_sfb,
SmemLayoutSFB{}(_,_,cute::Int<0>{}),
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
make_shape(shape<1>(TileShapeSFB{}), shape<2>(TileShapeSFB{})),
_1{}); // No programmatic multicast
return {
@@ -627,10 +638,14 @@ struct CollectiveMma<
// Partition the inputs based on the current block coordinates.
auto [m_coord, n_coord, k_coord, l_coord] = blk_coord;
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
Tensor gSFA = gSFA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
Tensor gSFB = gSFB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
auto broadcast_n = make_layout(
make_shape(Int<size<1>(TileShapeSFB{}) / size<1>(TileShape{})>{},
Int<cute::numeric_limits<int>::max()>{}),
make_stride(_0{}, size<1>(TileShapeSFB{}) / size<1>(TileShape{})));
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
Tensor gSFA = gSFA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
Tensor gSFB = gSFB_nkl(_,_,broadcast_n(n_coord),_,l_coord); // (BLK_N,BLK_K,k)
// Partition source and destination tensors for tma copies
Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k)
@@ -693,7 +708,8 @@ struct CollectiveMma<
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class FrgTensorC
class FrgTensorC,
class BlockCoord
>
CUTLASS_DEVICE void
mma(MainloopPipeline pipeline,
@@ -702,7 +718,8 @@ struct CollectiveMma<
int k_tile_count,
int thread_idx,
TensorStorage& shared_tensors,
[[maybe_unused]] Params const& params) {
[[maybe_unused]] Params const& params,
BlockCoord const& blk_coord) {
using namespace cute;
static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
@@ -712,7 +729,16 @@ struct CollectiveMma<
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
Tensor sSFA = make_tensor(make_smem_ptr(shared_tensors.smem_SFA.begin()), SmemLayoutSFA{}); // (BLK_M,BLK_K,PIPE)
Tensor sSFB = make_tensor(make_smem_ptr(shared_tensors.smem_SFB.begin()), SmemLayoutSFB{}); // (BLK_N,BLK_K,PIPE)
Tensor sSFB = [&]() {
if constexpr (size<1>(TileShape{}) >= 128) {
return make_tensor(make_smem_ptr(shared_tensors.smem_SFB.begin()), SmemLayoutSFB{}); // (BLK_N,BLK_K,PIPE)
}
else {
Tensor temp = make_tensor(make_smem_ptr(shared_tensors.smem_SFB.begin()), SmemLayoutSFB{}); // (BLK_SFB_N,BLK_K,PIPE)
auto n = get<1>(blk_coord);
return temp(make_coord(_,n % (size<1>(TileShapeSFB{}) / size<1>(TileShape{}))), _, _);
}
}();
//
// Define C accumulators and A/B partitioning
include/cutlass/gemm/collective/sm120_mma_array_tma_blockwise_scaling.hpp
+21 -2
View File
@@ -614,8 +614,6 @@ struct CollectiveMma<
}
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class FrgTensorC
>
@@ -627,6 +625,27 @@ struct CollectiveMma<
int thread_idx,
TensorStorage& shared_tensors,
Params const& mainloop_params) {
auto empty_tuple = make_tuple(_0{}, _0{}, _0{}, _0{});
mma(pipeline, smem_pipe_read, accum, k_tile_count,
thread_idx, shared_tensors, mainloop_params, empty_tuple);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class FrgTensorC,
class BlockCoord
>
CUTLASS_DEVICE void
mma(MainloopPipeline pipeline,
PipelineState smem_pipe_read,
FrgTensorC& accum,
int k_tile_count,
int thread_idx,
TensorStorage& shared_tensors,
Params const& mainloop_params,
[[maybe_unused]] BlockCoord& blk_crd) {
using namespace cute;
static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
include/cutlass/gemm/collective/sm120_mma_tma.hpp
+4 -2
View File
@@ -439,7 +439,8 @@ struct CollectiveMma<
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class FrgTensorC
class FrgTensorC,
class BlockCoord
>
CUTLASS_DEVICE void
mma(MainloopPipeline pipeline,
@@ -448,7 +449,8 @@ struct CollectiveMma<
int k_tile_count,
int thread_idx,
TensorStorage& shared_tensors,
Params const& mainloop_params) {
Params const& mainloop_params,
[[maybe_unused]] BlockCoord& blk_crd) {
using namespace cute;
static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
include/cutlass/gemm/collective/sm120_mma_tma_blockwise_scaling.hpp
+19 -2
View File
@@ -545,8 +545,6 @@ struct CollectiveMma<
}
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class FrgTensorC
>
@@ -558,6 +556,25 @@ struct CollectiveMma<
int thread_idx,
TensorStorage& shared_tensors,
Params const& mainloop_params) {
auto empty_tuple = make_tuple(_0{}, _0{}, _0{}, _0{});
mma(pipeline, smem_pipe_read, accum, k_tile_count, thread_idx, shared_tensors, mainloop_params, empty_tuple);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class FrgTensorC,
class BlockCoord
>
CUTLASS_DEVICE void
mma(MainloopPipeline pipeline,
PipelineState smem_pipe_read,
FrgTensorC& accum,
int k_tile_count,
int thread_idx,
TensorStorage& shared_tensors,
Params const& mainloop_params,
[[maybe_unused]] BlockCoord& blk_crd) {
using namespace cute;
static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp
+47 -11
View File
@@ -168,10 +168,28 @@ public:
static constexpr uint32_t NumProducerThreads = CollectiveMainloop::NumProducerThreadEvents;
static constexpr bool IsMainloopAuxiliaryLoadNeeded = detail::HasAuxiliaryLoad_v<typename CollectiveMainloop::DispatchPolicy>;
/// Register requirement for Load and Math WGs
static constexpr int RegsPerThread =
2 * size<0>(TileShape{}) * size<1>(TileShape{}) / NumMmaThreads *
sizeof(ElementAccumulator) / sizeof(uint32_t);
// Detect if this is SM120 blockscaled kernel which hits low register pressure
// on smaller tiles
template <typename T>
struct IsSm120BlockScaled : cute::false_type {};
template <int Stages, int SchedStages, class ClusterShape, class KernelSchedule>
struct IsSm120BlockScaled<MainloopSm120ArrayTmaWarpSpecializedBlockScaled<Stages, SchedStages, ClusterShape, KernelSchedule>>
: cute::true_type {};
static constexpr bool IsLowRegisterPressure = IsSm120BlockScaled<DispatchPolicy>::value && (RegsPerThread <= 64);
/// Register requirement for Load and Math WGs
static constexpr uint32_t LoadRegisterRequirement = 40;
static constexpr uint32_t MmaRegisterRequirement = 232;
static constexpr bool IsSm120Family = cute::is_same_v<typename DispatchPolicy::ArchTag, arch::Sm120>;
// 1 stage ordered sequence between mainloop and epilogue producer load threads
using LoadWarpOrderBarrier = cutlass::OrderedSequenceBarrier<1,2>;
@@ -606,7 +624,9 @@ public:
auto k_tile_count = size<3>(gA_mkl);
if (warp_group_role == WarpGroupRole::Producer) {
cutlass::arch::warpgroup_reg_dealloc<LoadRegisterRequirement>();
if constexpr (!IsLowRegisterPressure) {
cutlass::arch::warpgroup_reg_dealloc<LoadRegisterRequirement>();
}
if (producer_warp_role == ProducerWarpRole::Scheduler) {
// GroupScheduler requires a producer warp to iterate over the group infos and push
@@ -882,7 +902,9 @@ public:
} // Producer Warp Group End
else if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) {
cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
if constexpr (!IsLowRegisterPressure) {
cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
}
// Index of warp group within consumer warp groups
int consumer_warp_group_idx = warp_group_role == WarpGroupRole::Consumer0 ? 0 : 1;
@@ -935,15 +957,29 @@ public:
if (TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
accumulators,
work_k_tile_count,
mma_thread_idx,
shared_storage.tensors.mainloop,
params.mainloop
);
if constexpr (IsSm120Family) {
collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
accumulators,
work_k_tile_count,
mma_thread_idx,
shared_storage.tensors.mainloop,
params.mainloop,
blk_coord
);
}
else {
collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
accumulators,
work_k_tile_count,
mma_thread_idx,
shared_storage.tensors.mainloop,
params.mainloop
);
}
// Make sure the math instructions are done and free buffers before entering the epilogue
collective_mainloop.mma_tail(
include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp
+47 -11
View File
@@ -167,10 +167,28 @@ public:
static constexpr uint32_t NumProducerThreads = CollectiveMainloop::NumProducerThreadEvents;
static constexpr bool IsMainloopAuxiliaryLoadNeeded = detail::HasAuxiliaryLoad_v<typename CollectiveMainloop::DispatchPolicy>;
/// Register requirement for Load and Math WGs
static constexpr int RegsPerThread =
2 * size<0>(TileShape{}) * size<1>(TileShape{}) / (NumMmaWarpGroups * NumThreadsPerWarpGroup) *
sizeof(ElementAccumulator) / sizeof(uint32_t);
// Detect if this is SM120 blockscaled kernel which hits low register pressure
// on smaller tiles
template <typename T>
struct IsSm120BlockScaled : cute::false_type {};
template <int Stages, int SchedStages, class ClusterShape, class KernelSchedule>
struct IsSm120BlockScaled<MainloopSm120ArrayTmaWarpSpecializedBlockScaled<Stages, SchedStages, ClusterShape, KernelSchedule>>
: cute::true_type {};
static constexpr bool IsLowRegisterPressure = IsSm120BlockScaled<DispatchPolicy>::value && (RegsPerThread <= 64);
/// Register requirement for Load and Math WGs
static constexpr uint32_t LoadRegisterRequirement = 40;
static constexpr uint32_t MmaRegisterRequirement = 232;
static constexpr bool IsSm120Family = cute::is_same_v<typename DispatchPolicy::ArchTag, arch::Sm120>;
// 1 stage ordered sequence between mainloop and epilogue producer load threads
using LoadWarpOrderBarrier = cutlass::OrderedSequenceBarrier<1,2>;
@@ -644,7 +662,9 @@ public:
auto k_tile_count = size<3>(gA_mkl);
if (warp_group_role == WarpGroupRole::Producer) {
cutlass::arch::warpgroup_reg_dealloc<LoadRegisterRequirement>();
if constexpr (!IsLowRegisterPressure) {
cutlass::arch::warpgroup_reg_dealloc<LoadRegisterRequirement>();
}
if (producer_warp_role == ProducerWarpRole::Scheduler) {
// GroupScheduler requires a producer warp to iterate over the group infos and push
@@ -920,7 +940,9 @@ public:
} // Producer Warp Group End
else if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) {
cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
if constexpr (!IsLowRegisterPressure) {
cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
}
// Index of warp group within consumer warp groups
int consumer_warp_group_idx = warp_group_role == WarpGroupRole::Consumer0 ? 0 : 1;
@@ -975,15 +997,29 @@ public:
math_wg_order_barrier.wait();
collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
accumulators,
work_k_tile_count,
mma_thread_idx,
shared_storage.tensors.mainloop,
params.mainloop
);
if constexpr (IsSm120Family) {
collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
accumulators,
work_k_tile_count,
mma_thread_idx,
shared_storage.tensors.mainloop,
params.mainloop,
blk_coord
);
}
else {
collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
accumulators,
work_k_tile_count,
mma_thread_idx,
shared_storage.tensors.mainloop,
params.mainloop
);
}
math_wg_order_barrier.arrive();
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp
+23 -9
View File
@@ -786,15 +786,29 @@ public:
// MSVC CTAD breaks if we say "Tensor" here, so we use "auto" instead.
auto accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape)); // (MMA,MMA_M,MMA_N)
if (TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
accumulators,
work_k_tile_count,
mma_thread_idx,
shared_storage.tensors.mainloop,
params.mainloop
);
if constexpr (IsSm120Family) {
collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
accumulators,
work_k_tile_count,
mma_thread_idx,
shared_storage.tensors.mainloop,
params.mainloop,
blk_coord
);
}
else {
collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
accumulators,
work_k_tile_count,
mma_thread_idx,
shared_storage.tensors.mainloop,
params.mainloop
);
}
// Make sure the math instructions are done and free buffers before entering the epilogue
collective_mainloop.mma_tail(
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp
+23 -9
View File
@@ -849,15 +849,29 @@ public:
// Order two Math WG's MMA one after the other, helps hide Epilogue
math_wg_order_barrier.wait();
collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
accumulators,
k_tile_count,
warp_group_thread_idx,
shared_storage.tensors.mainloop,
params.mainloop
);
if constexpr (IsSm120Family) {
collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
accumulators,
k_tile_count,
warp_group_thread_idx,
shared_storage.tensors.mainloop,
params.mainloop,
blk_coord
);
}
else {
collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
accumulators,
k_tile_count,
warp_group_thread_idx,
shared_storage.tensors.mainloop,
params.mainloop
);
}
// Cue for next Math WG's MMA to start
math_wg_order_barrier.arrive();