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

Release v4.0.0 (#2294)

Browse Source
This commit is contained in:
Kihiro Bando
2025-05-13 15:55:29 -04:00
committed by GitHub
parent ad7b2f5e84
commit f115c3f854
299 changed files with 51495 additions and 4413 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/arch/config.hpp
+1
View File
@@ -93,6 +93,7 @@
# define CUTE_ARCH_STSM_SM90_ENABLED
#endif
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED))
# define CUTE_ARCH_TCGEN05_S8_MMA_ENABLED
#endif
include/cute/arch/mma_sm100_umma.hpp
+32
View File
@@ -1394,7 +1394,11 @@ struct SM100_MMA_MXF4_SS
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
"tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#else
"tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#endif
"}\n"
:
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC),
@@ -1411,7 +1415,11 @@ struct SM100_MMA_MXF4_SS
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
"tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#else
"tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#endif
"}\n"
:
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC),
@@ -1457,7 +1465,11 @@ struct SM100_MMA_MXF4NVF4_SS_SPARSE
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
"tcgen05.mma.sp.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, [%7], %3, [%5], [%6], p; \n\t"
#else
"tcgen05.mma.sp.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X [%0], %1, %2, [%7], %3, [%5], [%6], p; \n\t"
#endif
"}\n"
:
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC),
@@ -1475,7 +1487,11 @@ struct SM100_MMA_MXF4NVF4_SS_SPARSE
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
"tcgen05.mma.sp.cta_group::1.kind::mxf4.block_scale.block32 [%0], %1, %2, [%7], %3, [%5], [%6], p; \n\t"
#else
"tcgen05.mma.sp.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [%0], %1, %2, [%7], %3, [%5], [%6], p; \n\t"
#endif
"}\n"
:
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC),
@@ -1520,7 +1536,11 @@ struct SM100_MMA_MXF4_2x1SM_SS
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
"tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#else
"tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#endif
"}\n"
:
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC),
@@ -1537,7 +1557,11 @@ struct SM100_MMA_MXF4_2x1SM_SS
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
"tcgen05.mma.cta_group::2.kind::mxf4.block_scale.block32 [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#else
"tcgen05.mma.cta_group::2.kind::mxf4.block_scale.scale_vec::2X [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#endif
"}\n"
:
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC),
@@ -1582,7 +1606,11 @@ struct SM100_MMA_MXF4NVF4_2x1SM_SS_SPARSE
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
"tcgen05.mma.sp.cta_group::2.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, [%7], %3, [%5], [%6], p; \n\t"
#else
"tcgen05.mma.sp.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X [%0], %1, %2, [%7], %3, [%5], [%6], p; \n\t"
#endif
"}\n"
:
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC),
@@ -1600,7 +1628,11 @@ struct SM100_MMA_MXF4NVF4_2x1SM_SS_SPARSE
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
"tcgen05.mma.sp.cta_group::2.kind::mxf4.block_scale.block32 [%0], %1, %2, [%7], %3, [%5], [%6], p; \n\t"
#else
"tcgen05.mma.sp.cta_group::2.kind::mxf4.block_scale.scale_vec::2X [%0], %1, %2, [%7], %3, [%5], [%6], p; \n\t"
#endif
"}\n"
:
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC),
include/cute/atom/copy_traits_sm100_tma.hpp
+2 -2
View File
@@ -104,7 +104,7 @@ struct Copy_Traits<SM100_TMA_2SM_LOAD, NumBitsPerTMA, AuxParams_>
auto
get_tma_tensor(GShape const& g_shape) const {
static_assert(is_congruent<decltype(g_shape), decltype(aux_params_.g_stride_)>::value);
return make_counting_tensor(make_layout(g_shape, aux_params_.g_stride_));
return make_coord_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
// Don't try to execute a copy with SM100_TMA_2SM_LOAD before calling .with()
@@ -192,7 +192,7 @@ struct Copy_Traits<SM100_TMA_2SM_LOAD_MULTICAST, NumBitsPerTMA, AuxParams_>
auto
get_tma_tensor(GShape const& g_shape) const {
static_assert(is_congruent<decltype(g_shape), decltype(aux_params_.g_stride_)>::value);
return make_counting_tensor(make_layout(g_shape, aux_params_.g_stride_));
return make_coord_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
// Don't try to execute a copy with SM100_TMA_2SM_LOAD_MULTICAST_OP before calling .with()
include/cute/atom/copy_traits_sm90_tma.hpp
+4 -4
View File
@@ -146,7 +146,7 @@ struct Copy_Traits<SM90_TMA_LOAD, NumBitsPerTMA, AuxParams_>
auto
get_tma_tensor(GShape const& g_shape) const {
static_assert(is_congruent<decltype(g_shape), decltype(aux_params_.g_stride_)>::value);
return make_counting_tensor(make_layout(g_shape, aux_params_.g_stride_));
return make_coord_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
// Don't try to execute a copy with SM90_TMA_LOAD before calling .with()
@@ -276,7 +276,7 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST, NumBitsPerTMA, AuxParams_>
auto
get_tma_tensor(GShape const& g_shape) const {
static_assert(is_congruent<decltype(g_shape), decltype(aux_params_.g_stride_)>::value);
return make_counting_tensor(make_layout(g_shape, aux_params_.g_stride_));
return make_coord_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
// Don't try to execute a copy with SM90_TMA_LOAD_MULTICAST before calling .with()
@@ -350,7 +350,7 @@ struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, AuxParams_>
auto
get_tma_tensor(GShape const& g_shape) const {
static_assert(is_congruent<decltype(g_shape), decltype(aux_params_.g_stride_)>::value);
return make_counting_tensor(make_layout(g_shape, aux_params_.g_stride_));
return make_coord_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
// Construct new TMA_STORE with (unsafe) swapped out TMA descriptor ptr (for grouped gemm/ptr array gemm)
@@ -463,7 +463,7 @@ struct Copy_Traits<SM90_TMA_REDUCE_ADD, NumBitsPerTMA, AuxParams_>
auto
get_tma_tensor(GShape const& g_shape) const {
static_assert(is_congruent<decltype(g_shape), decltype(aux_params_.g_stride_)>::value);
return make_counting_tensor(make_layout(g_shape, aux_params_.g_stride_));
return make_coord_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
template <class Coord, int... Is>
include/cute/tensor_impl.hpp
+3 -3
View File
@@ -474,14 +474,14 @@ make_fragment_like(Tensor<Engine,Layout> const& tensor)
}
//
// make_counting_tensor
// make_coord_tensor
// Make a tensor from a layout by binding it to a counting iter with 0-offset of the same profile as the codomain.
//
template <class Layout, __CUTE_REQUIRES(is_layout<Layout>::value)>
CUTE_HOST_DEVICE constexpr
auto
make_counting_tensor(Layout const& layout)
make_coord_tensor(Layout const& layout)
{
return make_tensor(make_inttuple_iter(coprofile(layout)), layout);
}
@@ -496,7 +496,7 @@ CUTE_HOST_DEVICE constexpr
auto
make_identity_tensor(Shape const& shape)
{
return make_counting_tensor(make_identity_layout(shape));
return make_coord_tensor(make_identity_layout(shape));
}
//
include/cutlass/arch/config.h
+4 -4
View File
@@ -105,10 +105,8 @@
/////////////////////////////////////////////////////////////////////////////////////////////////
// SM101 and SM101a
#if !CUTLASS_CLANG_CUDA && (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 8))
#if !CUTLASS_CLANG_CUDA && (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 8)
#define CUTLASS_ARCH_MMA_SM101_SUPPORTED 1
#if (!defined(CUTLASS_ARCH_MMA_SM101_ENABLED) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1010)
#define CUTLASS_ARCH_MMA_SM101_ENABLED 1
@@ -118,7 +116,7 @@
#endif
// SM101f
#if (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9))
#if !CUTLASS_CLANG_CUDA && (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
#define CUTLASS_ARCH_MMA_SM101F_SUPPORTED 1
#endif
@@ -130,6 +128,8 @@
/////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////
// SM120 and SM120a
#if !CUTLASS_CLANG_CUDA && (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 8))
#define CUTLASS_ARCH_MMA_SM120_SUPPORTED 1
include/cutlass/epilogue/collective/builders/sm120_builder.inl
+40
View File
@@ -284,6 +284,46 @@ struct CallbacksBuilder<
>;
};
// Overload CallbacksBuilder to pick the correct copy atoms for PtrArray epilogue fusions
template <
int StagesC,
int StagesD,
int FragmentSize,
bool ReuseSmemC,
bool DelayTmaStore,
int NumEpilogueWarpgroups,
class FusionOp,
class TileShape_MNK,
class EpilogueTile_MN,
class AccLoadOp,
class ElementAccumulator
>
struct CallbacksBuilder<
Sm120PtrArrayTmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore, NumEpilogueWarpgroups>,
FusionOp,
TileShape_MNK,
EpilogueTile_MN,
ElementAccumulator,
AccLoadOp,
cute::enable_if_t<(FusionOp::IsAuxOutSupported ^ FusionOp::IsAuxInSupported) // only one aux tensor
&& not cute::is_subbyte_v<typename FusionOp::ElementAux>>
> {
using GmemStrideTypeAux = gemm::TagToStrideC_t<typename FusionOp::GmemLayoutTagAux>;
using SmemLayoutAtomAux = decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom<
GmemStrideTypeAux, typename FusionOp::ElementAux, EpilogueTile_MN>());
using CopyOpR2S = decltype(detail::sm120_get_smem_store_op_for_accumulator<GmemStrideTypeAux, typename FusionOp::ElementAux>());
using CopyOpS2R = decltype(detail::sm120_get_smem_load_op_for_source<GmemStrideTypeAux, typename FusionOp::ElementAux>());
using SmemCopyOpAux = cute::conditional_t<FusionOp::IsAuxOutSupported, CopyOpR2S, CopyOpS2R>;
using Callbacks = fusion::FusionCallbacks<
Sm120PtrArrayTmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore, NumEpilogueWarpgroups>,
FusionOp, TileShape_MNK, EpilogueTile_MN,
SmemLayoutAtomAux, SmemCopyOpAux
>;
};
// Helper for building TMA warp-specialized collective epilogues, specialized by
// the fusion operation performed and the dispatch policy to use.
include/cutlass/epilogue/collective/builders/sm90_builder.inl
+35 -14
View File
@@ -115,12 +115,13 @@ sm90_compute_tile_shape_or_override() {
if constexpr (cute::is_same_v<EpilogueTileType, EpilogueTileAuto>) {
auto epi_tile = [&] () {
if constexpr (detail::sm90_is_cooperative_v<Schedule>) {
auto tile_m = cute::min(_128{}, size<0>(TileShape_MNK{}));
auto tile_n = cute::gcd(cute::min(_32{}, size<1>(TileShape_MNK{})), size<1>(TileShape_MNK{}));
return make_shape(tile_m, tile_n);
}
else if constexpr (detail::sm90_is_warp_specialized_v<Schedule>) {
constexpr int N_perf = sizeof_bits_v<ElementD> == 8 ? 64 : 32;
constexpr int N_perf = (sizeof_bits_v<ElementD> == 8) && (size<1>(TileShape_MNK{}) % 64 == 0) ? 64 : 32;
auto tile_m = cute::min(_64{}, size<0>(TileShape_MNK{}));
auto tile_n = cute::gcd(cute::min(Int<N_perf>{}, size<1>(TileShape_MNK{})), size<1>(TileShape_MNK{}));
return make_shape(tile_m, tile_n);
@@ -194,9 +195,9 @@ struct CallbacksBuilder<
using SmemLayoutAtomAux = decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom<
GmemStrideTypeAux, typename FusionOp::ElementAux, EpilogueTile_MN>());
using CopyOpR2S = decltype(detail::sm90_get_smem_store_op_for_accumulator<
GmemStrideTypeAux, typename FusionOp::ElementAux>());
GmemStrideTypeAux, typename FusionOp::ElementAux, EpilogueTile_MN>());
using CopyOpS2R = decltype(detail::sm90_get_smem_load_op_for_source<
GmemStrideTypeAux, typename FusionOp::ElementAux>());
GmemStrideTypeAux, typename FusionOp::ElementAux, EpilogueTile_MN>());
using SmemCopyOpAux = cute::conditional_t<FusionOp::IsAuxOutSupported, CopyOpR2S, CopyOpS2R>;
using Callbacks = fusion::FusionCallbacks<
@@ -234,9 +235,9 @@ struct CallbacksBuilder<
using SmemLayoutAtomAux = decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom<
GmemStrideTypeAux, typename FusionOp::ElementAux, EpilogueTile_MN>());
using CopyOpR2S = decltype(detail::sm90_get_smem_store_op_for_accumulator<
GmemStrideTypeAux, typename FusionOp::ElementAux>());
GmemStrideTypeAux, typename FusionOp::ElementAux, EpilogueTile_MN>());
using CopyOpS2R = decltype(detail::sm90_get_smem_load_op_for_source<
GmemStrideTypeAux, typename FusionOp::ElementAux>());
GmemStrideTypeAux, typename FusionOp::ElementAux, EpilogueTile_MN>());
using SmemCopyOpAux = cute::conditional_t<FusionOp::IsAuxOutSupported, CopyOpR2S, CopyOpS2R>;
using Callbacks = fusion::FusionCallbacks<
@@ -316,7 +317,17 @@ struct Sm90TmaBuilderImpl {
>;
// Get the smallest tiled copy we can use to retile the accumulators
using CopyAtomC = Copy_Atom<SM90_U32x4_STSM_N, cutlass::half_t>;
// using CopyAtomC = Copy_Atom<SM90_U32x4_STSM_N, cutlass::half_t>;
using CopyAtomC = cute::conditional_t<
size<1>(EpilogueTile_MN{}) % 16 == 0,
Copy_Atom<SM90_U32x4_STSM_N, cutlass::half_t>,
cute::conditional_t<
size<1>(EpilogueTile_MN{}) % 8 == 0,
Copy_Atom<SM90_U32x2_STSM_N, cutlass::half_t>,
void
>
>;
static_assert(!cute::is_same_v<CopyAtomC, void>, "CopyAtomC can't be void, divisiblity check for EpilogueTile_MN failed");
// Get register to register tiled copy that happen before shared memory store.
// Apply void as no register transform op needed currently.
using CopyOpR2R = void;
@@ -343,10 +354,10 @@ struct Sm90TmaBuilderImpl {
FusionCallbacks,
CopyOpG2S,
decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom<UnderlyingGmemStrideTypeC, ElementC, EpilogueTile_MN>()),
decltype(detail::sm90_get_smem_load_op_for_source<UnderlyingGmemStrideTypeC, ElementC>()),
decltype(detail::sm90_get_smem_load_op_for_source<UnderlyingGmemStrideTypeC, ElementC, EpilogueTile_MN>()),
CopyOpS2G,
decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom<UnderlyingGmemStrideTypeD, ElementD, EpilogueTile_MN>()),
decltype(detail::sm90_get_smem_store_op_for_accumulator<UnderlyingGmemStrideTypeD, ElementD>()),
decltype(detail::sm90_get_smem_store_op_for_accumulator<UnderlyingGmemStrideTypeD, ElementD, EpilogueTile_MN>()),
CopyAtomC,
CopyOpR2R
>;
@@ -404,7 +415,7 @@ struct AuxLoadDescriptor {
>()
);
using CopyOpS2R =
decltype(detail::sm90_get_smem_load_op_for_source<Stride, ElementAux>());
decltype(detail::sm90_get_smem_load_op_for_source<Stride, ElementAux, EpilogueTile>());
};
// Get Stride, SmemLayout, and CopyOpS2R for AuxStore node
@@ -425,7 +436,7 @@ struct AuxStoreDescriptor {
>()
);
using CopyOpR2S =
decltype(detail::sm90_get_smem_store_op_for_accumulator<Stride, ElementAux>());
decltype(detail::sm90_get_smem_store_op_for_accumulator<Stride, ElementAux, EpilogueTile>());
};
} // namespace detail
@@ -745,7 +756,7 @@ private:
using SmemLayoutAtomAux = decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom<
GmemStrideTypeAux, typename Schedule::ElementT, EpilogueTile_MN>());
using SmemCopyOpAux = decltype(detail::sm90_get_smem_store_op_for_accumulator<
GmemStrideTypeAux, typename Schedule::ElementT>());
GmemStrideTypeAux, typename Schedule::ElementT, EpilogueTile_MN>());
using FusionOperationAux = fusion::LinCombPerRowBiasEltActAux<
GmemLayoutTagD, Schedule::template ActivationFunctor, ElementD, ElementCompute,
typename Schedule::ElementT, typename Schedule::ElementBias, ElementC_, ElementCompute
@@ -769,7 +780,17 @@ private:
using GmemStrideTypeD = gemm::TagToStrideC_t<GmemLayoutTagD>;
// Get the smallest tiled copy we can use to retile the accumulators
using CopyAtomC = Copy_Atom<SM90_U32x4_STSM_N, cutlass::half_t>;
using CopyAtomC = cute::conditional_t<
size<1>(EpilogueTile_MN{}) % 16 == 0,
Copy_Atom<SM90_U32x4_STSM_N, cutlass::half_t>,
cute::conditional_t<
size<1>(EpilogueTile_MN{}) % 8 == 0,
Copy_Atom<SM90_U32x2_STSM_N, cutlass::half_t>,
void
>
>;
static_assert(!cute::is_same_v<CopyAtomC, void>, "CopyAtomC can't be void, divisiblity check for EpilogueTile_MN failed");
// Get register to register tiled copy that happen before shared memory store.
// Apply void as no register transform op needed.
using CopyOpR2R = void;
@@ -788,10 +809,10 @@ public:
cute::conditional_t<Schedule::StoreT, FusionCallbacksAux, FusionCallbacksNoAux>,
SM90_TMA_LOAD,
decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom<GmemStrideTypeC, ElementC, EpilogueTile_MN>()),
decltype(detail::sm90_get_smem_load_op_for_source<GmemStrideTypeC, ElementC>()),
decltype(detail::sm90_get_smem_load_op_for_source<GmemStrideTypeC, ElementC, EpilogueTile_MN>()),
SM90_TMA_STORE,
decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom<GmemStrideTypeD, ElementD, EpilogueTile_MN>()),
decltype(detail::sm90_get_smem_store_op_for_accumulator<GmemStrideTypeD, ElementD>()),
decltype(detail::sm90_get_smem_store_op_for_accumulator<GmemStrideTypeD, ElementD, EpilogueTile_MN>()),
CopyAtomC,
CopyOpR2R
>;
include/cutlass/epilogue/collective/builders/sm90_common.inl
+21 -5
View File
@@ -37,16 +37,26 @@ namespace cutlass::epilogue::collective::detail {
///////////////////////////////////////////////////////////////////////////////
// Selects the largest vectorized smem store atom available
template <class GmemStrideTypeD, class ElementD>
template <class GmemStrideTypeD, class ElementD, class EpilogueTile_MN>
constexpr auto
sm90_get_smem_store_op_for_accumulator() {
using namespace cute;
if constexpr (sizeof(ElementD) == 2 && size<0>(GmemStrideTypeD{}) == 1) {
return SM90_U16x8_STSM_T{};
if constexpr (size<1>(EpilogueTile_MN{}) % 16 == 0) {
return SM90_U16x8_STSM_T{};
}
else if constexpr (size<1>(EpilogueTile_MN{}) % 8 == 0) {
return SM90_U16x4_STSM_T{};
}
}
else if constexpr (sizeof(ElementD) == 2 && size<1>(GmemStrideTypeD{}) == 1) {
return SM90_U32x4_STSM_N{};
if constexpr (size<1>(EpilogueTile_MN{}) % 16 == 0) {
return SM90_U32x4_STSM_N{};
}
else if constexpr (size<1>(EpilogueTile_MN{}) % 8 == 0) {
return SM90_U32x2_STSM_N{};
}
}
else {
// auto-vectorizing store
@@ -55,20 +65,26 @@ sm90_get_smem_store_op_for_accumulator() {
}
// Selects the largest vectorized smem load atom available
template <class GmemStrideTypeC, class ElementC>
template <class GmemStrideTypeC, class ElementC, class EpilogueTile_MN>
constexpr auto
sm90_get_smem_load_op_for_source() {
using namespace cute;
// Reuse the logic from smem store selector
using SmemStoreOp = decltype(sm90_get_smem_store_op_for_accumulator<GmemStrideTypeC, ElementC>());
using SmemStoreOp = decltype(sm90_get_smem_store_op_for_accumulator<GmemStrideTypeC, ElementC, EpilogueTile_MN>());
if constexpr (cute::is_same_v<SmemStoreOp, SM90_U16x8_STSM_T>) {
return SM75_U16x8_LDSM_T{};
}
else if constexpr (cute::is_same_v<SmemStoreOp, SM90_U16x4_STSM_T>) {
return SM75_U16x4_LDSM_T{};
}
else if constexpr (cute::is_same_v<SmemStoreOp, SM90_U32x4_STSM_N>) {
return SM75_U32x4_LDSM_N{};
}
else if constexpr (cute::is_same_v<SmemStoreOp, SM90_U32x2_STSM_N>) {
return SM75_U32x2_LDSM_N{};
}
else {
// auto-vectorizing load
return AutoVectorizingCopyWithAssumedAlignment<128>{};
include/cutlass/epilogue/collective/default_epilogue.hpp
+2 -2
View File
@@ -215,8 +215,8 @@ public:
Tensor cD_mn = local_tile(mD_crd, take<0,2>(blk_shape_MNK), make_coord(m_coord, n_coord)); // (BLK_M,BLK_N)
Tensor tCcD_mn = thr_mma.partition_C(cD_mn); // (VEC,THR_M,THR_N)
// Relative coordinate tensors (static)
Tensor cD = make_counting_tensor(cD_mn.layout()); // (BLK_M,BLK_N)
Tensor tCcD = make_counting_tensor(tCcD_mn.layout()); // (VEC,THR_M,THR_N)
Tensor cD = make_coord_tensor(cD_mn.layout()); // (BLK_M,BLK_N)
Tensor tCcD = make_coord_tensor(tCcD_mn.layout()); // (VEC,THR_M,THR_N)
// Subtract the global "bottom right" corner from the local "top left" corner to get the max relative coordinate
auto residue_cD = shape_MN - cD_mn(_0{}); // (m,n)
auto residue_tCcD = shape_MN - tCcD_mn(_0{}); // (m,n)
include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp
+8 -12
View File
@@ -286,12 +286,8 @@ public:
void* workspace) {
// These tensor shapes (only applicable for grouped gemm) and pointers are only used to create tensormap/tma desc.
// These will be replaced with correct values before the initial tma load.
auto init_shape = repeat_like(append<4>(typename ProblemShape::UnderlyingProblemShape{}, 1), int32_t(1));
// These tensor shapes (only applicable for grouped gemm) and pointers are only used to create tensormap/tma desc.
// These will be replaced with correct values before the initial tma load.
constexpr int tma_alignment_bits = 128;
auto init_M = tma_alignment_bits;
auto init_N = tma_alignment_bits;
auto init_M = int32_t(size<0>(CtaTileShape{}));
auto init_N = int32_t(size<1>(CtaTileShape{}));
auto init_L = 1;
InternalStrideC stride_c;
@@ -745,8 +741,8 @@ public:
Tensor cD_mn = local_tile(mD_crd, take<0,2>(cta_tile_mnk), make_coord(m_coord, n_coord)); // (CTA_M,CTA_N)
Tensor tTR_cD_mn = thread_t2r.partition_D(flat_divide(cD_mn, EpilogueTile{})); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
// Relative coordinate tensors (static)
Tensor cD = make_counting_tensor(cD_mn.layout()); // (CTA_M,CTA_N)
Tensor tTR_cD = make_counting_tensor(tTR_cD_mn.layout()); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
Tensor cD = make_coord_tensor(cD_mn.layout()); // (CTA_M,CTA_N)
Tensor tTR_cD = make_coord_tensor(tTR_cD_mn.layout()); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
// Subtract the global "bottom right" corner from the local "top left" corner to get the max relative coordinate
auto residue_cD = make_coord(M,N) - cD_mn(_0{}); // (m,n)
auto residue_tTR_cD = make_coord(M,N) - tTR_cD_mn(_0{}); // (m,n)
@@ -786,8 +782,8 @@ public:
[[maybe_unused]] bool reverse_epi_n = ReuseTmem && acc_pipe_consumer_state.phase() == 0;
static_assert(not (ReuseTmem && AccumulatorPipeline::Stages != 1), "Tmem reuse requires 1 accumulator stage");
// Predication for TMA store (one warp issues TMA store)
bool issue_tma_store = warp_idx == 0;
// Predication for TMA store (a single thread from one warp issues TMA store)
bool issue_tma_store = (warp_idx == 0) && cute::elect_one_sync();
// In the reuse smem configuration we have StagesC smem buffers and at most StagesD committed TMA stores in flight.
// The TMA store pipeline producer acquire returns when at most StagesD-1 committed stores are in-flight, so we can
@@ -1118,8 +1114,8 @@ public:
Tensor cD_mn = local_tile(mD_crd, take<0,2>(cta_tile_mnk), make_coord(m_coord, n_coord)); // (CTA_M,CTA_N)
Tensor tTR_cD_mn = thread_t2r.partition_D(flat_divide(cD_mn, EpilogueTile{})); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
// Relative coordinate tensors (static)
Tensor cD = make_counting_tensor(cD_mn.layout()); // (CTA_M,CTA_N)
Tensor tTR_cD = make_counting_tensor(tTR_cD_mn.layout()); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
Tensor cD = make_coord_tensor(cD_mn.layout()); // (CTA_M,CTA_N)
Tensor tTR_cD = make_coord_tensor(tTR_cD_mn.layout()); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
// Subtract the global "bottom right" corner from the local "top left" corner to get the max relative coordinate
auto residue_cD = make_coord(M,N) - cD_mn(_0{}); // (m,n)
auto residue_tTR_cD = make_coord(M,N) - tTR_cD_mn(_0{}); // (m,n)
include/cutlass/epilogue/collective/sm100_epilogue_nosmem.hpp
+2 -2
View File
@@ -606,8 +606,8 @@ public:
// Construct the EVT consumer callbacks
auto residue_cD = make_coord(M,N) - cD(_0{});
auto residue_tTR_cD = make_coord(M,N) - tTR_cD(_0{});
Tensor cD_ = make_counting_tensor(cD.layout());
Tensor tTR_cD_ = make_counting_tensor(tTR_cD.layout());
Tensor cD_ = make_coord_tensor(cD.layout());
Tensor tTR_cD_ = make_coord_tensor(tTR_cD.layout());
constexpr bool RefSrc = false;
Tensor mC = make_tensor(make_gmem_ptr<GmemElementC>(params.ptr_C), make_shape(M,N,L), params.dC);
include/cutlass/epilogue/collective/sm100_epilogue_tma_warpspecialized.hpp
+4 -4
View File
@@ -695,8 +695,8 @@ public:
Tensor cD_mn = local_tile(mD_crd, take<0,2>(cta_tile_mnk), make_coord(m_coord, n_coord)); // (CTA_M,CTA_N)
Tensor tTR_cD_mn = thread_t2r.partition_D(flat_divide(cD_mn, EpilogueTile{})); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
// Relative coordinate tensors (static)
Tensor cD = make_counting_tensor(cD_mn.layout()); // (CTA_M,CTA_N)
Tensor tTR_cD = make_counting_tensor(tTR_cD_mn.layout()); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
Tensor cD = make_coord_tensor(cD_mn.layout()); // (CTA_M,CTA_N)
Tensor tTR_cD = make_coord_tensor(tTR_cD_mn.layout()); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
// Subtract the global "bottom right" corner from the local "top left" corner to get the max relative coordinate
auto residue_cD = make_coord(M,N) - cD_mn(_0{}); // (m,n)
auto residue_tTR_cD = make_coord(M,N) - tTR_cD_mn(_0{}); // (m,n)
@@ -1065,8 +1065,8 @@ public:
Tensor cD_mn = local_tile(mD_crd, take<0,2>(cta_tile_mnk), make_coord(m_coord, n_coord)); // (CTA_M,CTA_N)
Tensor tTR_cD_mn = thread_t2r.partition_D(flat_divide(cD_mn, EpilogueTile{})); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
// Relative coordinate tensors (static)
Tensor cD = make_counting_tensor(cD_mn.layout()); // (CTA_M,CTA_N)
Tensor tTR_cD = make_counting_tensor(tTR_cD_mn.layout()); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
Tensor cD = make_coord_tensor(cD_mn.layout()); // (CTA_M,CTA_N)
Tensor tTR_cD = make_coord_tensor(tTR_cD_mn.layout()); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
// Subtract the global "bottom right" corner from the local "top left" corner to get the max relative coordinate
auto residue_cD = make_coord(M,N) - cD_mn(_0{}); // (m,n)
auto residue_tTR_cD = make_coord(M,N) - tTR_cD_mn(_0{}); // (m,n)
include/cutlass/epilogue/collective/sm90_epilogue_array_tma_warpspecialized.hpp
+7 -11
View File
@@ -304,11 +304,9 @@ public:
[[maybe_unused]] void* workspace) {
// These tensor shapes (only applicable for grouped gemm) and pointers are only used to create tensormap/tma desc.
// These will be replaced with correct values before the initial tma load.
auto init_shape = repeat_like(append<4>(typename ProblemShape::UnderlyingProblemShape{}, 1), int32_t(1));
constexpr int tma_alignment_bits = 128;
auto init_M = tma_alignment_bits;
auto init_N = tma_alignment_bits;
auto init_L = get<3>(init_shape);
auto init_M = int32_t(size<0>(CtaTileMNK{}));
auto init_N = int32_t(size<1>(CtaTileMNK{}));
auto init_L = 1;
static_assert(!is_im2col_C and !is_im2col_D, "Im2Col not supported on C or D");
@@ -324,8 +322,6 @@ public:
auto problem_shape_MNKL = append<4>(problem_shape.get_host_problem_shape(0), 1);
init_M = get<0>(problem_shape_MNKL);
init_N = get<1>(problem_shape_MNKL);
init_L = get<3>(problem_shape_MNKL);
stride_c = args.dC;
stride_d = args.dD;
}
@@ -755,8 +751,8 @@ public:
Tensor cD_mn = local_tile(mD_crd, take<0,2>(CtaTileMNK{}), make_coord(m_coord, n_coord)); // (CTA_M,CTA_N)
Tensor tRS_cD_mn = thread_r2s.partition_S(flat_divide(cD_mn, EpilogueTile{})); // (R2S,R2S_M,R2S_N,EPI_M,EPI_N)
// Relative coordinate tensors (static)
Tensor cD = make_counting_tensor(cD_mn.layout()); // (CTA_M,CTA_N)
Tensor tRS_cD = make_counting_tensor(tRS_cD_mn.layout()); // (R2S,R2S_M,R2S_N,EPI_M,EPI_N)
Tensor cD = make_coord_tensor(cD_mn.layout()); // (CTA_M,CTA_N)
Tensor tRS_cD = make_coord_tensor(tRS_cD_mn.layout()); // (R2S,R2S_M,R2S_N,EPI_M,EPI_N)
// Subtract the global "bottom right" corner from the local "top left" corner to get the max relative coordinate
auto residue_cD = make_coord(M,N) - cD_mn(_0{}); // (m,n)
auto residue_tRS_cD = make_coord(M,N) - tRS_cD_mn(_0{}); // (m,n)
@@ -803,8 +799,8 @@ public:
// to ensure visibility of smem reads/writes to threads or TMA unit
auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(size(TiledMma{}), cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); };
// Predication for TMA store (one warp issues TMA store)
bool issue_tma_store = (thread_idx / NumThreadsPerWarp) == 0;
// Predication for TMA store (a single thread from one warp issues TMA store)
bool issue_tma_store = ((thread_idx / NumThreadsPerWarp) == 0) && cute::elect_one_sync();
// In the reuse smem configuration we have StagesC smem buffers and at most StagesD committed TMA stores in flight.
// The TMA store pipeline producer acquire returns when at most StagesD-1 committed stores are in-flight, so we can
include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp
+2 -2
View File
@@ -662,8 +662,8 @@ public:
}
}();
// Relative coordinate tensors (static)
Tensor cD = make_counting_tensor(cD_mn.layout()); // (CTA_M,CTA_N)
Tensor tRS_cD = make_counting_tensor(tRS_cD_mn.layout()); // (R2S,R2S_M,R2S_N,EPI_M,EPI_N)
Tensor cD = make_coord_tensor(cD_mn.layout()); // (CTA_M,CTA_N)
Tensor tRS_cD = make_coord_tensor(tRS_cD_mn.layout()); // (R2S,R2S_M,R2S_N,EPI_M,EPI_N)
// Subtract the global "bottom right" corner from the local "top left" corner to get the max relative coordinate
auto residue_cD = make_coord(M,N) - cD_mn(_0{}); // (m,n)
auto residue_tRS_cD = make_coord(M,N) - tRS_cD_mn(_0{}); // (m,n)
include/cutlass/epilogue/thread/activation.h
+17 -17
View File
@@ -496,13 +496,29 @@ template <>
struct HardSwish<float> {
using T = float;
static const bool kIsHeavy = false;
static constexpr float kOneSixth = 0.16666667f;
CUTLASS_HOST_DEVICE
T operator()(T const &x) const {
minimum<T> mn;
maximum<T> mx;
T relu6 = mn(mx(x + T(3), T(0)), T(6));
return x * relu6 * 0.16666667f;
return x * relu6 * kOneSixth;
}
};
template <>
struct HardSwish<cutlass::half_t> {
using T = cutlass::half_t;
static const bool kIsHeavy = false;
static constexpr float kOneSixth = 0.16666667f;
CUTLASS_HOST_DEVICE
T operator()(T const &x) const {
minimum<T> mn;
maximum<T> mx;
T relu6 = mn(mx(x + T(3), T(0)), T(6));
return x * relu6 * T(kOneSixth);
}
};
@@ -524,22 +540,6 @@ struct HardSwish<Array<T, N> > {
}
};
template <int N>
struct HardSwish<Array<half_t, N> > {
using T = half_t;
static const bool kIsHeavy = false;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value) const {
minimum<Array<T, N> > mn;
maximum<Array<T, N> > mx;
multiplies<Array<T, N> > mul;
plus<Array<T, N> > add;
return mul(mul(mn(mx(add(value, T(3)), T(0)), T(6)), value), T(0.16666667f));
}
};
template <typename T>
using ScaledHardSwish = Scale<HardSwish<T>>;
include/cutlass/functional.h
+1 -1
View File
@@ -722,7 +722,7 @@ struct has_unqualified_conj : cutlass::platform::false_type
template<typename T>
struct has_unqualified_conj<
T,
decltype(conj(cutlass::platform::declval<T>()), void())
decltype(static_cast<void>(conj(cutlass::platform::declval<T>())), void())
> : cutlass::platform::true_type
{};
include/cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl
+10 -6
View File
@@ -129,7 +129,10 @@ auto sm100_make_simt_gmem_tiled_copy_SFA() {
using ScaleCopyTypeA = cute::uint_byte_t<Alignment>;
using SmemScalingCopyAtomA = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ScaleCopyTypeA>, Element>;
constexpr int ElementsPerSFACopy = static_cast<int>(sizeof(ScaleCopyTypeA) / sizeof(Element));
return make_tiled_copy(SmemScalingCopyAtomA{}, Layout<Shape<_32>>{}, Layout<Shape<Int<ElementsPerSFACopy>>>{});
return make_tiled_copy(
SmemScalingCopyAtomA{},
Layout<Shape<_32>>{}, // 32 threads
Layout<Shape<Shape<Int<ScaleGranularityM>, Int<ElementsPerSFACopy>>>, Stride<Stride<_0, _1>>>{});
}
else {
using SmemScalingCopyAtomA = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<Element>, Element>;
@@ -138,9 +141,8 @@ auto sm100_make_simt_gmem_tiled_copy_SFA() {
}
else {
// we expect scale Ks per tile to be small
constexpr int LeadingScalesPerTileSFA = ScaleKsPerTile;
using SmemScalingCopyAtomA = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<Element>, Element>;
return make_tiled_copy(SmemScalingCopyAtomA{}, Layout<Shape<_1, Int<LeadingScalesPerTileSFA>>>{}, Layout<Shape<_1,_1>>{});
return make_tiled_copy(SmemScalingCopyAtomA{}, Layout<Shape<_32>>{}, Layout<Shape<_1>>{});
}
}
@@ -161,7 +163,10 @@ auto sm100_make_simt_gmem_tiled_copy_SFB() {
using ScaleCopyTypeB = cute::uint_byte_t<Alignment>;
using SmemScalingCopyAtomB = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ScaleCopyTypeB>, Element>;
constexpr int ElementsPerSFBCopy = static_cast<int>(sizeof(ScaleCopyTypeB) / sizeof(Element));
return make_tiled_copy(SmemScalingCopyAtomB{}, Layout<Shape<_32>>{}, Layout<Shape<Int<ElementsPerSFBCopy>>>{});
return make_tiled_copy(
SmemScalingCopyAtomB{},
Layout<Shape<_32>>{}, // 32 threads
Layout<Shape<Shape<Int<ScaleGranularityN>, Int<ElementsPerSFBCopy>>>, Stride<Stride<_0, _1>>>{});
}
else {
using SmemScalingCopyAtomB = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<Element>, Element>;
@@ -170,9 +175,8 @@ auto sm100_make_simt_gmem_tiled_copy_SFB() {
}
else {
// we expect scale Ks per tile to be small
constexpr int LeadingScalesPerTileSFB = ScaleKsPerTile;
using SmemScalingCopyAtomB = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<Element>, Element>;
return make_tiled_copy(SmemScalingCopyAtomB{}, Layout<Shape<_1, Int<LeadingScalesPerTileSFB>>>{}, Layout<Shape<_1,_1>>{});
return make_tiled_copy(SmemScalingCopyAtomB{}, Layout<Shape<_32>>{}, Layout<Shape<_1>>{});
}
}
include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl
+12 -5
View File
@@ -260,7 +260,9 @@ struct CollectiveBuilder<
GmmaMajorB, ElementBMma, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
static constexpr size_t TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 2 /* for A and B */ : 0;
static constexpr int KernelSmemCarveout = static_cast<int>(TensorMapStorage);
static constexpr size_t SchedulerPipelineStorage = cute::is_pointer_v<TagToStrideA_t<GmemLayoutATag>> ?
sizeof(cutlass::PipelineDetail::PipelineAsyncSharedStorage<8>) : 0;
static constexpr int KernelSmemCarveout = static_cast<int>(TensorMapStorage + SchedulerPipelineStorage);
static constexpr int Sm90ReducedSmemCapacityBytes = detail::sm90_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int PipelineStages = detail::compute_stage_count_or_override<Sm90ReducedSmemCapacityBytes,
@@ -446,7 +448,9 @@ public:
// Handle mixed dtype array GEMM's size of tensor map storage.
static constexpr size_t TensorMapStorage = sizeof(cute::TmaDescriptor) * size_t(IsMixedInput) * 4;
static constexpr int KernelSmemCarveout = static_cast<int>(TensorMapStorage);
static constexpr size_t SchedulerPipelineStorage = cute::is_pointer_v<TagToStrideA_t<GmemLayoutATag_>> ?
sizeof(cutlass::PipelineDetail::PipelineAsyncSharedStorage<8>) : 0;
static constexpr int KernelSmemCarveout = static_cast<int>(TensorMapStorage + SchedulerPipelineStorage);
static constexpr int Sm90ReducedSmemCapacityBytes = detail::sm90_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int PipelineStages = IsMixedInput ?
@@ -570,7 +574,9 @@ struct CollectiveBuilder<
GmmaMajorB, ElementB, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
static constexpr size_t TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 2 /* for A and B */ : 0;
static constexpr int KernelSmemCarveout = static_cast<int>(TensorMapStorage);
static constexpr size_t SchedulerPipelineStorage = cute::is_pointer_v<TagToStrideA_t<GmemLayoutATag>> ?
sizeof(cutlass::PipelineDetail::PipelineAsyncSharedStorage<8>) : 0;
static constexpr int KernelSmemCarveout = static_cast<int>(TensorMapStorage + SchedulerPipelineStorage);
static constexpr int Sm90ReducedSmemCapacityBytes = detail::sm90_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int PipelineStages = detail::compute_stage_count_or_override<Sm90ReducedSmemCapacityBytes,
@@ -1128,8 +1134,9 @@ struct CollectiveBuilder<
static constexpr size_t TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 2 /* for A and B */ : 0;
// Reserve 128B for 8 stages of tile scheduling
static constexpr size_t TileSchedulerCarveout = IsArrayOfPointersGemm ? 128 : 0;
static constexpr int KernelSmemCarveout = static_cast<int>(TensorMapStorage) + static_cast<int>(TileSchedulerCarveout);
static constexpr size_t SchedulerPipelineStorage = cute::is_pointer_v<TagToStrideA_t<GmemLayoutATag>> ?
sizeof(cutlass::PipelineDetail::PipelineAsyncSharedStorage<8>) : 0;
static constexpr int KernelSmemCarveout = static_cast<int>(TensorMapStorage + SchedulerPipelineStorage);
static constexpr int ScaleMsPerTile = size<0>(TileShape_MNK{}) / ScaleGranularityM;
static constexpr int ScaleNsPerTile = size<1>(TileShape_MNK{}) / ScaleGranularityN;
include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp
+3 -4
View File
@@ -427,10 +427,9 @@ struct CollectiveMma<
cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) {
// These tensor shapes (only applicable for grouped gemm) and pointers are only used to create tensormap/tma desc.
// These will be replaced with correct values before the initial tma load.
constexpr int tma_alignment_bits = 128;
auto init_M = tma_alignment_bits;
auto init_N = tma_alignment_bits;
auto init_K = tma_alignment_bits;
auto init_M = int32_t(size<0>(TileShape{}));
auto init_N = int32_t(size<1>(TileShape{}));
auto init_K = int32_t(size<2>(TileShape{}));
auto init_L = 1;
// Tensor pointers will be fixed before the first access
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_blockwise_scaling.hpp
+10 -4
View File
@@ -190,8 +190,14 @@ struct CollectiveMma<
using TransformB = TransformB_;
using ArchTag = typename DispatchPolicy::ArchTag;
static constexpr int AlignmentSFA = GmemTiledCopySFA::AtomNumVal::value * sizeof(typename GmemTiledCopySFA::ValType) / sizeof(ElementAccumulator);
static constexpr int AlignmentSFB = GmemTiledCopySFB::AtomNumVal::value * sizeof(typename GmemTiledCopySFB::ValType) / sizeof(ElementAccumulator);
static constexpr int CopyAlignmentSFA = GmemTiledCopySFA::AtomNumVal::value * sizeof(typename GmemTiledCopySFA::ValType) / sizeof(ElementAccumulator);
static constexpr int CopyAlignmentSFB = GmemTiledCopySFB::AtomNumVal::value * sizeof(typename GmemTiledCopySFB::ValType) / sizeof(ElementAccumulator);
static constexpr int AlignmentSFA = CopyAlignmentSFA * (GmemTiledCopySFA::AtomNumVal::value > 1 ?
(size<0,1>(InternalLayoutSFA{}.stride()) == 1 ? ScaleGranularityM : ScaleGranularityK) : 1);
static constexpr int AlignmentSFB = CopyAlignmentSFB * (GmemTiledCopySFB::AtomNumVal::value > 1 ?
(size<0,1>(InternalLayoutSFB{}.stride()) == 1 ? ScaleGranularityN : ScaleGranularityK) : 1);
using MainloopABPipeline = cutlass::PipelineTmaUmmaAsync<
DispatchPolicy::Stages,
@@ -522,8 +528,8 @@ struct CollectiveMma<
auto [M,N,K,L] = problem_shape_MNKL;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), InternalStrideA{});
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), InternalStrideB{});
implementable_sf = implementable_sf && cutlass::detail::check_alignment<AlignmentSFA>(ScaleConfig::tile_atom_to_shape_SFA(problem_shape_MNKL));
implementable_sf = implementable_sf && cutlass::detail::check_alignment<AlignmentSFB>(ScaleConfig::tile_atom_to_shape_SFB(problem_shape_MNKL));
implementable_sf = implementable_sf && cutlass::detail::check_alignment<CopyAlignmentSFA>(ScaleConfig::tile_atom_to_shape_SFA(problem_shape_MNKL));
implementable_sf = implementable_sf && cutlass::detail::check_alignment<CopyAlignmentSFB>(ScaleConfig::tile_atom_to_shape_SFB(problem_shape_MNKL));
if (!implementable_sf) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for Scale Factors.\n");
}
include/cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp
+10 -4
View File
@@ -201,8 +201,14 @@ struct CollectiveMma<
AtomThrShapeMNK>;
using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState;
static constexpr int AlignmentSFA = GmemTiledCopySFA::AtomNumVal::value * sizeof(typename GmemTiledCopySFA::ValType) / sizeof(ElementAccumulator);
static constexpr int AlignmentSFB = GmemTiledCopySFB::AtomNumVal::value * sizeof(typename GmemTiledCopySFB::ValType) / sizeof(ElementAccumulator);
static constexpr int CopyAlignmentSFA = GmemTiledCopySFA::AtomNumVal::value * sizeof(typename GmemTiledCopySFA::ValType) / sizeof(ElementAccumulator);
static constexpr int CopyAlignmentSFB = GmemTiledCopySFB::AtomNumVal::value * sizeof(typename GmemTiledCopySFB::ValType) / sizeof(ElementAccumulator);
static constexpr int AlignmentSFA = CopyAlignmentSFA * (GmemTiledCopySFA::AtomNumVal::value > 1 ?
(size<0,1>(LayoutSFA{}.stride()) == 1 ? ScaleGranularityM : ScaleGranularityK) : 1);
static constexpr int AlignmentSFB = CopyAlignmentSFB * (GmemTiledCopySFB::AtomNumVal::value > 1 ?
(size<0,1>(LayoutSFB{}.stride()) == 1 ? ScaleGranularityN : ScaleGranularityK) : 1);
// Two arrivals per thread in the warp (1 arrival and 1 arrival through cp.async.mbarrier)
static constexpr int NumMainloopSFProducerThreadEvents = 64;
@@ -566,8 +572,8 @@ struct CollectiveMma<
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
}
bool implementable_sf = cutlass::detail::check_alignment<AlignmentSFA>(args.layout_SFA);
implementable_sf = implementable_sf && cutlass::detail::check_alignment<AlignmentSFB>(args.layout_SFB);
bool implementable_sf = cutlass::detail::check_alignment<CopyAlignmentSFA>(args.layout_SFA);
implementable_sf = implementable_sf && cutlass::detail::check_alignment<CopyAlignmentSFB>(args.layout_SFB);
if (!implementable_sf) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for Scale Factors.\n");
include/cutlass/gemm/collective/sm100_mma_warpspecialized_mixed_input.hpp
-824
View File
@@ -1,824 +0,0 @@
/***************************************************************************************************
* Copyright (c) 2023 - 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.
*
**************************************************************************************************/
#pragma once
#include <cuda_bf16.h>
#include "cutlass/cutlass.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/numeric_conversion.h"
#include "cutlass/detail/sm100_tmem_helper.hpp"
#include "cutlass/detail/cluster.hpp"
#include "cute/algorithm/functional.hpp"
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/atom/copy_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/arch/mma_sm100.hpp"
#include "cutlass/trace.h"
#include "cutlass/kernel_hardware_info.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
using namespace cute;
/////////////////////////////////////////////////////////////////////////////////////////////////
// WarpSpecialized Mainloop for Mixed Input Kernels
template <
int Load2TransformPipelineStageCount_,
int Transform2MmaPipelineStageCount_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape,
class TileShape_,
class ElementA_,
class StrideA_,
class ElementB_,
class StrideB_,
class TiledMma_,
class GmemTiledCopyA_,
class SmemLayoutAtomsA_,
class CopyAtomsA_,
class TransformA_,
class GmemTiledCopyB_,
class SmemLayoutAtomsB_,
class CopyAtomsB_,
class TransformB_>
struct CollectiveMma<
MainloopSm100TmaUmmaWarpSpecializedMixedInput<
Load2TransformPipelineStageCount_,
Transform2MmaPipelineStageCount_,
SchedulerPipelineStageCount_,
AccumulatorPipelineStageCount_,
ClusterShape>,
TileShape_,
ElementA_,
StrideA_,
ElementB_,
StrideB_,
TiledMma_,
GmemTiledCopyA_,
SmemLayoutAtomsA_,
CopyAtomsA_,
TransformA_,
GmemTiledCopyB_,
SmemLayoutAtomsB_,
CopyAtomsB_,
TransformB_>
{
//
// Type Aliases
//
// Determine MMA type: MMA_1SM vs MMA_2SM
using AtomThrShapeMNK = Shape<decltype(shape<0>(typename TiledMma_::ThrLayoutVMNK{})), _1, _1>;
using DispatchPolicy = MainloopSm100TmaUmmaWarpSpecializedMixedInput<
Load2TransformPipelineStageCount_,
Transform2MmaPipelineStageCount_,
SchedulerPipelineStageCount_,
AccumulatorPipelineStageCount_,
ClusterShape>;
using TileShape = TileShape_;
using TiledMma = TiledMma_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
// Define A and B block shapes for reduced size TMA_LOADs
using CtaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{}))));
using CtaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{}))));
using ElementA = ElementA_;
using StrideA = StrideA_;
using ElementAMma = typename TiledMma::ValTypeA;
static constexpr int IsSubbyteA = cute::sizeof_bits_v<ElementA> < 8;
using TmaElementA = cute::conditional_t<IsSubbyteA, uint8_t, ElementA>;
using ElementB = ElementB_;
using StrideB = StrideB_;
using ElementBMma = typename TiledMma::ValTypeB;
using StrideScale = cute::Stride<cute::Int<1>, int64_t, int64_t>;
using NonVoidStrideScale = cute::conditional_t<
cute::is_void_v<StrideScale>, cute::Stride<_1, int64_t, int64_t>, StrideScale>;
using ElementAccumulator = typename TiledMma::ValTypeC;
using GmemTiledCopyA = GmemTiledCopyA_;
using GmemTiledCopyB = GmemTiledCopyB_;
using SmemLayoutAtomsA = SmemLayoutAtomsA_;
using SmemLayoutAtomsB = SmemLayoutAtomsB_;
using CopyAtomsA = CopyAtomsA_;
using CopyAtomsB = CopyAtomsB_;
using TransformA = TransformA_;
using TransformB = TransformB_;
using ArchTag = typename DispatchPolicy::ArchTag;
static_assert(sizeof(ElementA) < 2, "Matrix to be scaled should be provided in A otherwise input is not supported");
static_assert(cute::is_same_v<ElementAMma, cutlass::bfloat16_t> || cute::is_same_v<ElementAMma, cutlass::half_t> || cute::is_same_v<ElementAMma, cutlass::float_e4m3_t>, "Compute type A should be cutlass::bfloat16_t or cutlass::half_t or cutlass::float_e4m3_t");
using Load2TransformPipeline = cutlass::PipelineTmaTransformAsync<
DispatchPolicy::Load2TransformPipelineStageCount,
AtomThrShapeMNK>;
using Load2TransformPipelineState = typename Load2TransformPipeline::PipelineState;
using Transform2MmaPipeline = cutlass::PipelineUmmaConsumerAsync<
DispatchPolicy::Transform2MmaPipelineStageCount,
AtomThrShapeMNK>;
using Transform2MmaPipelineState = typename Transform2MmaPipeline::PipelineState;
using Mma2AccumPipeline = cutlass::PipelineUmmaAsync<
DispatchPolicy::Schedule::AccumulatorPipelineStageCount,
AtomThrShapeMNK>;
using Mma2AccumPipelineState = typename Mma2AccumPipeline::PipelineState;
// Thread Counts
static constexpr uint32_t NumAccumThreads = 128; //Maintains compatibility with input_transform kernel
static constexpr uint32_t NumTransformationThreads = 128;
// Get the Algorithm parameters
constexpr static int AccumulatorPipelineStageCount = DispatchPolicy::Schedule::AccumulatorPipelineStageCount;
constexpr static int StagesPerTile = size<2>(CtaShapeA_MK{});
using SmemLayoutAtomA = typename SmemLayoutAtomsA::InputLayoutAtom;
using SmemLayoutAtomACompute = typename SmemLayoutAtomsA::ComputeLayoutAtom;
using SmemLayoutAtomB = typename SmemLayoutAtomsB::InputLayoutAtom;
using SmemLayoutAtomBCompute = typename SmemLayoutAtomsB::ComputeLayoutAtom;
using InputCopyAtomA = typename CopyAtomsA::InputCopyAtom;
using ComputeCopyAtomA = typename CopyAtomsA::ComputeCopyAtom;
using InputCopyAtomB = typename CopyAtomsB::InputCopyAtom;
using ComputeCopyAtomB = typename CopyAtomsB::ComputeCopyAtom;
static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
static_assert(((size<0,0>(CtaShapeA_MK{}) * size<1>(CtaShapeA_MK{})) % size<0>(SmemLayoutAtomACompute{})) == 0, "SmemLayoutAtomCompute must evenly divide tile shape.");
static_assert(((size<0,1>(CtaShapeA_MK{}) * size<2>(CtaShapeA_MK{})) % size<1>(SmemLayoutAtomACompute{})) == 0, "SmemLayoutAtomCompute must evenly divide tile shape.");
static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
static_assert(((size<0,0>(CtaShapeB_NK{}) * size<1>(CtaShapeB_NK{})) % size<0>(SmemLayoutAtomBCompute{})) == 0, "SmemLayoutAtomCompute must evenly divide tile shape.");
static_assert(((size<0,1>(CtaShapeB_NK{}) * size<2>(CtaShapeB_NK{})) % size<1>(SmemLayoutAtomBCompute{})) == 0, "SmemLayoutAtomCompute must evenly divide tile shape.");
// Tile along K mode first before tiling over MN. PIPE mode last as usual.
// This maximizes TMA boxes due to better smem-K vectorization, reducing total issued TMAs.
using SmemLayoutA = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomA{},
append(CtaShapeA_MK{}, Int<DispatchPolicy::Load2TransformPipelineStageCount>{}),
(cute::conditional_t<cutlass::gemm::detail::is_mn_major<StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})));
using SmemLayoutACompute = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomACompute{},
append(CtaShapeA_MK{}, Int<DispatchPolicy::Load2TransformPipelineStageCount>{}),
(cute::conditional_t<cutlass::gemm::detail::is_mn_major<StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})));
using SmemLayoutB = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomB{},
append(CtaShapeB_NK{}, Int<DispatchPolicy::Load2TransformPipelineStageCount>{}),
(cute::conditional_t<cutlass::gemm::detail::is_mn_major<StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})));
static_assert(DispatchPolicy::Load2TransformPipelineStageCount >= 2 && DispatchPolicy::Load2TransformPipelineStageCount >= 2,
"Specialization requires Stages set to value 2 or more.");
static_assert((cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value ||
cute::is_base_of<cute::UMMA::tmem_frg_base, typename TiledMma::FrgTypeA>::value ) &&
cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
"MMA atom must A operand from SMEM or TMEM and B operand from SMEM for this mainloop.");
static_assert((cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>),
"GmemTiledCopyA - invalid TMA copy atom specified.");
static_assert((cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>),
"GmemTiledCopyB - invalid TMA copy atom specified.");
struct PipelineStorage {
using Load2TransformPipelineStorage = typename Load2TransformPipeline::SharedStorage;
alignas(16) Load2TransformPipelineStorage load2transform_pipeline;
using Transform2MmaPipelineStorage = typename Transform2MmaPipeline::SharedStorage;
alignas(16) Transform2MmaPipelineStorage transform2mma_pipeline;
using Mma2AccumPipelineStorage = typename Mma2AccumPipeline::SharedStorage;
alignas(16) Mma2AccumPipelineStorage mma2accum_pipeline;
};
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128, _0> {
struct TensorStorageUntransformed {
cute::ArrayEngine<ElementA, cute::cosize_v<SmemLayoutA>> smem_A;
cute::ArrayEngine<ElementB, cute::cosize_v<SmemLayoutB>> smem_B;
};
struct TensorStorageTransformedAinSmem {
alignas(1024) cute::ArrayEngine<ElementAMma, cute::cosize_v<SmemLayoutACompute>> smem_ACompute;
alignas(1024) cute::ArrayEngine<ElementBMma, cute::cosize_v<SmemLayoutB>> smem_BCompute;
};
union TensorStorageTransformedAinTmem {
alignas(1024) cute::ArrayEngine<ElementAMma, 1> smem_ACompute; // No smem_ACompute
alignas(1024) cute::ArrayEngine<ElementBMma, cute::cosize_v<SmemLayoutB>> smem_BCompute;
};
using TensorStorageTransformed = cute::conditional_t<
cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value,
TensorStorageTransformedAinSmem,
TensorStorageTransformedAinTmem>;
TensorStorageUntransformed input;
TensorStorageTransformed compute;
} tensors;
PipelineStorage pipeline;
};
using TensorStorage = typename SharedStorage::TensorStorage;
// Different from other GEMM kernels, both CTAs should be aware of loads. Both CTAs will work on
// loaded input A and B matrices to convert the data type
static constexpr uint32_t TmaTransactionBytes =
cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * size<2>(SmemLayoutA{}) * static_cast<uint32_t>(sizeof_bits<ElementA>::value))+
cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * size<2>(SmemLayoutB{}) * static_cast<uint32_t>(sizeof_bits<ElementB>::value));
// Host side kernel arguments
struct Arguments {
ElementA const* ptr_A{nullptr};
StrideA dA{};
ElementB const* ptr_B{nullptr};
StrideB dB{};
};
// Device side kernel params
struct Params {
using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return<IsDynamicCluster>(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})),
make_tile(typename TiledMma::AtomThrID{})));
using TMA_A = decltype(make_tma_atom_A_sm100<TmaElementA>(
GmemTiledCopyA{},
make_tensor(static_cast<ElementA const*>(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}),
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
ClusterLayout_VMNK{})
);
using TMA_B = decltype(make_tma_atom_B_sm100<ElementB>(
GmemTiledCopyB{},
make_tensor(static_cast<ElementB const*>(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}),
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
ClusterLayout_VMNK{})
);
TMA_A tma_load_a;
TMA_B tma_load_b;
TMA_A tma_load_a_fallback;
TMA_B tma_load_b_fallback;
dim3 cluster_shape_fallback;
};
CUTLASS_DEVICE
CollectiveMma(Params const& params, ClusterShape cluster_shape, uint32_t block_rank_in_cluster)
: cluster_shape_(cluster_shape)
, block_rank_in_cluster_(block_rank_in_cluster) {
if constexpr (IsDynamicCluster) {
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y);
observed_tma_load_a_ = is_fallback_cluster ? &params.tma_load_a_fallback : &params.tma_load_a;
observed_tma_load_b_ = is_fallback_cluster ? &params.tma_load_b_fallback : &params.tma_load_b;
}
else {
observed_tma_load_a_ = &params.tma_load_a;
observed_tma_load_b_ = &params.tma_load_b;
}
}
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) {
(void) workspace;
// Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
Tensor tensor_a = make_tensor(args.ptr_A, make_layout(make_shape(M,K,L), args.dA));
Tensor tensor_b = make_tensor(args.ptr_B, make_layout(make_shape(N,K,L), args.dB));
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape);
// Cluster layout for TMA construction
auto cluster_layout_vmnk = tiled_divide(make_layout(cluster_shape), make_tile(typename TiledMma::AtomThrID{}));
auto cluster_shape_fallback = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape_fallback);
// Cluster layout for TMA construction
auto cluster_layout_vmnk_fallback = tiled_divide(make_layout(cluster_shape_fallback), make_tile(typename TiledMma::AtomThrID{}));
typename Params::TMA_A tma_load_a = make_tma_atom_A_sm100<TmaElementA>(
GmemTiledCopyA{},
tensor_a,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_B tma_load_b = make_tma_atom_B_sm100<ElementB>(
GmemTiledCopyB{},
tensor_b,
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_A tma_load_a_fallback = make_tma_atom_A_sm100<TmaElementA>(
GmemTiledCopyA{},
tensor_a,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk_fallback);
typename Params::TMA_B tma_load_b_fallback = make_tma_atom_B_sm100<ElementB>(
GmemTiledCopyB{},
tensor_b,
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk_fallback);
return {
tma_load_a,
tma_load_b,
tma_load_a_fallback,
tma_load_b_fallback,
hw_info.cluster_shape_fallback
};
}
template<class ProblemShape>
static bool
can_implement(
ProblemShape const& problem_shape,
[[maybe_unused]] Arguments const& args) {
constexpr int tma_alignment_bits = 128;
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
bool implementable = true;
constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits<ElementB>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{});
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
}
return implementable;
}
/// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
CUTLASS_DEVICE static void
prefetch_tma_descriptors(Params const& params) {
if constexpr (IsDynamicCluster) {
dim3 cs = cute::cluster_shape();
const bool is_fallback_cluster = (cs.x == params.cluster_shape_fallback.x && cs.y == params.cluster_shape_fallback.y);
if (is_fallback_cluster) {
cute::prefetch_tma_descriptor(params.tma_load_a_fallback.get_tma_descriptor());
cute::prefetch_tma_descriptor(params.tma_load_b_fallback.get_tma_descriptor());
}
else {
cute::prefetch_tma_descriptor(params.tma_load_a.get_tma_descriptor());
cute::prefetch_tma_descriptor(params.tma_load_b.get_tma_descriptor());
}
}
else {
cute::prefetch_tma_descriptor(params.tma_load_a.get_tma_descriptor());
cute::prefetch_tma_descriptor(params.tma_load_b.get_tma_descriptor());
}
}
/// Construct A Single Stage's Accumulator Shape
CUTLASS_DEVICE auto
partition_accumulator_shape() {
auto acc_shape = partition_shape_C(TiledMma{}, take<0,2>(TileShape{})); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N)
return acc_shape;
}
/// Produce the inputs to the transform threads by loading inputs from gmem -> smem
template <
class GTensorA, class GTensorB,
class GTensorPartitionedA, class GTensorPartitionedB,
class STensorA, class STensorB,
class TileCoordMNKL,
class KTileIterator
>
CUTLASS_DEVICE auto
load(
Params const& params,
Load2TransformPipeline pipeline,
Load2TransformPipelineState load2xform_pipeline_state,
cute::tuple<GTensorA, GTensorB,
GTensorPartitionedA, GTensorPartitionedB,
STensorA, STensorB,
uint16_t, uint16_t> const& load_inputs,
TileCoordMNKL const& cta_coord_mnkl,
KTileIterator k_tile_iter, int k_tile_count) {
auto [unused_gA, unused_gB,
tAgA_mkl, tBgB_nkl, tAsA, tBsB,
mcast_mask_a, mcast_mask_b] = load_inputs;
// slice out the work coord from tiled tensors
Tensor tAgA = tAgA_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl));
Tensor tBgB = tBgB_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
uint32_t skip_wait = (k_tile_count <= 0);
auto pipeline_flag = pipeline.producer_try_acquire(load2xform_pipeline_state, skip_wait);
// Issue the Mainloop loads
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 0; --k_tile_count) {
// LOCK mainloop_load2xform_pipeline_state for _writing_
pipeline.producer_acquire(load2xform_pipeline_state, pipeline_flag);
int write_stage = load2xform_pipeline_state.index();
using BarrierType = typename Load2TransformPipeline::ProducerBarrierType;
BarrierType* tma_barrier = pipeline.producer_get_barrier(load2xform_pipeline_state);
// Advance mainloop_pipe
++load2xform_pipeline_state;
skip_wait = (k_tile_count <= 1);
pipeline_flag = pipeline.producer_try_acquire(load2xform_pipeline_state, skip_wait);
copy(observed_tma_load_a_->with(*tma_barrier, mcast_mask_a), tAgA(_,*k_tile_iter), tAsA(_,write_stage));
copy(observed_tma_load_b_->with(*tma_barrier, mcast_mask_b), tBgB(_,*k_tile_iter), tBsB(_,write_stage));
++k_tile_iter;
}
return cute::make_tuple(load2xform_pipeline_state, k_tile_iter);
}
/// Set up the data needed by this collective for load.
/// Returned tuple must contain at least two elements, with the first two elements being:
/// gA_mkl - The tiled tensor for input A
/// gB_nkl - The tiled tensor for input B
// Other inputs needed for load(): partitioned AB tensors for gmem and smem, and mcast masks
template <class ProblemShape_MNKL>
CUTLASS_DEVICE auto
load_init(
ProblemShape_MNKL const& problem_shape_MNKL,
Params const& params,
TensorStorage& shared_storage) const {
auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL);
ThrMMA cta_mma = TiledMma{}.get_slice(blockIdx.x % size(typename TiledMma::AtomThrID{}));
Tensor tCgA_mkl = cta_mma.partition_A(gA_mkl); // (MMA, MMA_M, MMA_K, m, k, l)
Tensor tCgB_nkl = cta_mma.partition_B(gB_nkl); // (MMA, MMA_N, MMA_K, n, k, l)
Tensor sA = make_tensor(make_smem_ptr(shared_storage.input.smem_A.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_storage.input.smem_B.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE)
// Define the CTA-in-cluster Layout and Coord
Layout cta_layout_mnk = make_layout(cluster_shape_);
Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMma::AtomThrID{}));
auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster_);
// Project the cta_layout for tma_a along the n-modes
auto [tAgA_mkl, tAsA] = tma_partition(*observed_tma_load_a_,
get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),
group_modes<0,3>(sA), group_modes<0,3>(tCgA_mkl));
// Project the cta_layout for tma_b along the m-modes
auto [tBgB_nkl, tBsB] = tma_partition(*observed_tma_load_b_,
get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),
group_modes<0,3>(sB), group_modes<0,3>(tCgB_nkl));
// TMA Multicast Masks
uint16_t mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);
uint16_t mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);
return cute::make_tuple(
gA_mkl, gB_nkl, // for scheduler
tAgA_mkl, tBgB_nkl, tAsA, tBsB, // for input tensor values
mcast_mask_a, mcast_mask_b); // multicast masks
}
template<
class KTileIterator, class Accumulator,
class GTensorA, class DstCopyA, class SrcTensorA, class DstTensorA,
class GTensorB
>
CUTLASS_DEVICE auto
transform(
Load2TransformPipeline load2transform_pipeline,
Load2TransformPipelineState load2transform_pipeline_consumer_state,
Transform2MmaPipeline transform2mma_pipeline,
Transform2MmaPipelineState transform2mma_pipeline_producer_state,
Accumulator accumulators,
cute::tuple<GTensorA, DstCopyA, SrcTensorA, DstTensorA,
GTensorB> input_operands,
KTileIterator k_tile_iter, int k_tile_count) {
cutlass::arch::NamedBarrier transform_bar(NumTransformationThreads, cutlass::arch::ReservedNamedBarriers::TransformBarrier);
// tAsA : (Copy,#Copy),MMA_Rest,MMA_M_Rest,MMA_K_Rest, SmemStages (In SMEM)
// tAdA : (Copy,#Copy),MMA_Rest,MMA_M_Rest,MMA_K_Rest, NumComputeMtxs, SmemStages (In SMEM or TMEM)
// tBsB : (Copy,#Copy),MMA_Rest,MMA_N_Rest,MMA_K_Rest, SmemStages (In SMEM)
// tBsB : (Copy,#Copy),MMA_Rest,MMA_N_Rest,MMA_K_Rest, NumComputeMtxs, SmemStages (In SMEM)
auto [unused_tAgA, dst_copy_A, tAsA, tAsACompute,
unused_tBgB] = input_operands;
// Create the tensors in registers
auto tArA = make_tensor<ElementA>(tAsA(_,_,_,_,0).shape());
auto tArACompute = make_tensor<ElementAMma>(tAsA(_,_,_,_,0).shape());
auto tArA_x2 = recast<Array<ElementA,2>>(tArA);
auto tArACompute_x2 = recast<Array<ElementAMma,2>>(tArACompute);
uint32_t skip_wait = (k_tile_count <= 0);
auto load2transform_flag = load2transform_pipeline.consumer_try_wait(load2transform_pipeline_consumer_state, skip_wait);
auto transform2mma_flag = transform2mma_pipeline.producer_try_acquire(transform2mma_pipeline_producer_state, skip_wait);
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 0; --k_tile_count) {
load2transform_pipeline.consumer_wait(load2transform_pipeline_consumer_state, load2transform_flag);
transform2mma_pipeline.producer_acquire(transform2mma_pipeline_producer_state, transform2mma_flag);
int load2transform_consumer_index = load2transform_pipeline_consumer_state.index(); // read stage
int transform2mma_producer_index = transform2mma_pipeline_producer_state.index(); //write stage
auto curr_load2transform_pipeline_consumer_state = load2transform_pipeline_consumer_state;
auto curr_transform2mma_pipeline_producer_state = transform2mma_pipeline_producer_state;
// Copy the input A matrix from SMEM
copy(AutoVectorizingCopy{}, tAsA(_,_,_,_,load2transform_consumer_index), tArA);
//Transform Input A stored in registers
cute::transform(tArA_x2, tArACompute_x2, cutlass::NumericArrayConverter<ElementAMma, ElementA, 2, cutlass::FloatRoundStyle::round_to_nearest_satfinite>::convert);
//Transformed A stored in TMEM
copy(dst_copy_A, tArACompute, tAsACompute(_,_,_,_,transform2mma_producer_index));
// Loads from SMEM are done. Signal the mainloop load as early as possible
transform_bar.sync();
load2transform_pipeline.consumer_release(curr_load2transform_pipeline_consumer_state);
// fence for SMEM writes
cutlass::arch::fence_view_async_shared();
if constexpr (is_tmem<decltype(tAsACompute)>::value) {
// fence for TMEM writes if A operand is coming from TMEM
cutlass::arch::fence_view_async_tmem_store();
}
// Let the MMA know we are done transforming
transform2mma_pipeline.producer_commit(curr_transform2mma_pipeline_producer_state);
// Next pipeline stage
++load2transform_pipeline_consumer_state;
++transform2mma_pipeline_producer_state;
skip_wait = (k_tile_count <= 1);
// Peek the next pipeline stage's barriers
load2transform_flag = load2transform_pipeline.consumer_try_wait(load2transform_pipeline_consumer_state, skip_wait);
transform2mma_flag = transform2mma_pipeline.producer_try_acquire(transform2mma_pipeline_producer_state, skip_wait);
}
return cute::make_tuple(load2transform_pipeline_consumer_state, transform2mma_pipeline_producer_state);
}
template<class ProblemShape_MNKL, class Accumulator>
CUTLASS_DEVICE auto
transform_init(
Params const& params,
ProblemShape_MNKL const& problem_shape_MNKL,
Accumulator accumulators,
TensorStorage& shared_storage) {
auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL);
Tensor sA_orig = make_tensor(make_smem_ptr(shared_storage.input.smem_A.begin()), SmemLayoutA{});
Tensor sA = as_position_independent_swizzle_tensor(sA_orig); //tCsA
Tensor sACompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), SmemLayoutACompute{}); //tCsACompute
// Map input, compute, and fragment tensors to
// Copy strategies and partitioned tensors. These will become the input
// operands of the transform function. Depending on MMA atom type, the
// operands can reside in SMEM or TMEM
auto setup_copy_ops = [&] (
auto tensor_input,
auto input_copy_atom,
auto tensor_compute,
auto make_fragment,
auto compute_copy_atom) constexpr {
auto fragment_compute = make_fragment(tensor_compute); //tCrA(Compute)
if constexpr (cute::is_tmem<cute::remove_cvref_t<decltype(fragment_compute)>>::value) {
// For M=128 with 2CTA MMA atoms, the TMEM tensor for A has a duplicated allocation.
// Instead of allocation a 64x16 TMEM tensor, we have a 128x16 allocation
// See: TmemAllocMode::Duplicated.
Tensor tensor_input2x = [&] () constexpr {
if constexpr (decltype(size<0,0>(fragment_compute) == Int<128>{} && size<0,0>(tensor_input) == Int<64>{})::value) {
return make_tensor(tensor_input.data(),
logical_product(tensor_input.layout(),
make_tile(make_tile(Layout<_2,_0>{},_),_,_,_))); // ((128,16),m,k,PIPE)
}
else {
return tensor_input;
}
}(); //tCsA_2x
fragment_compute.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators); //tCrA.data()
auto reg2tmem_tiled_copy = make_tmem_copy(compute_copy_atom, fragment_compute(_,_,0,0));
auto thr_reg2tmem_tiled_copy = reg2tmem_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads);
auto partitioned_tensor_input = thr_reg2tmem_tiled_copy.partition_S(tensor_input2x);
auto partitioned_tensor_compute = thr_reg2tmem_tiled_copy.partition_D(fragment_compute);
return cute::make_tuple(reg2tmem_tiled_copy, partitioned_tensor_input, partitioned_tensor_compute);
}
else {
auto tensor_compute_ind_sw = as_position_independent_swizzle_tensor(tensor_compute);
auto reg2smem_tiled_copy = make_cotiled_copy(compute_copy_atom, Layout<Shape <_128,_8>, Stride< _8,_1>>{},
tensor_compute(_,_,0,0).layout());
auto thr_reg2smem_tiled_copy = reg2smem_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads);
auto partitioned_tensor_input = thr_reg2smem_tiled_copy.partition_S(tensor_input);
auto partitioned_tensor_compute = thr_reg2smem_tiled_copy.partition_D(tensor_compute_ind_sw);
return cute::make_tuple(AutoVectorizingCopy{}, partitioned_tensor_input, partitioned_tensor_compute);
}
};
auto [dst_copy_A, tAsA, tAsACompute] =
setup_copy_ops(sA, InputCopyAtomA{}, sACompute, [&](auto &arg) {return TiledMma::make_fragment_A(arg);}, ComputeCopyAtomA{});
return cute::make_tuple(gA_mkl, dst_copy_A, tAsA, tAsACompute,
gB_nkl);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class FrgEngine, class FrgLayout,
class TensorA, class TensorB
>
CUTLASS_DEVICE auto
mma(
Transform2MmaPipeline transform2mma_pipeline,
Transform2MmaPipelineState transform2mma_pipeline_consumer_state,
Mma2AccumPipeline mma2accum_pipeline,
Mma2AccumPipelineState mma2accum_pipeline_producer_state,
cute::Tensor<FrgEngine, FrgLayout> const& accumulators,
cute::tuple<TensorA, TensorB> const& input_operands,
int k_tile_count
) {
TiledMma tiled_mma;
auto curr_transform2mma_pipeline_consumer_state = transform2mma_pipeline_consumer_state;
auto next_transform2mma_pipeline_consumer_state = transform2mma_pipeline_consumer_state;
uint32_t skip_wait = (k_tile_count <= 0);
auto transform2mma_flag = transform2mma_pipeline.consumer_try_wait(next_transform2mma_pipeline_consumer_state, skip_wait);
++next_transform2mma_pipeline_consumer_state;
// tCrA : (MMA), MMA_M, MMA_K, SmemStage (In SMEM or TMEM)
// We use SMEM stages to match #buffers in Load <-> Convert
// tCrB : (MMA), MMA_N, MMA_K, SmemStages (In SMEM)
auto const [tCrA, tCrB] = input_operands;
int remaining_accum_promotions = k_tile_count;
uint32_t mma2accum_skip_wait = (remaining_accum_promotions <= 0);
auto mma2accum_flag = mma2accum_pipeline.producer_try_acquire(mma2accum_pipeline_producer_state, mma2accum_skip_wait);
mma2accum_pipeline.producer_acquire(mma2accum_pipeline_producer_state, mma2accum_flag);
auto curr_mma2accum_pipeline_producer_state = mma2accum_pipeline_producer_state;
++mma2accum_pipeline_producer_state;
// No accumulator addition to the k_tile initially
tiled_mma.accumulate_ = UMMA::ScaleOut::Zero;
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 0; --k_tile_count) {
transform2mma_pipeline.consumer_wait(curr_transform2mma_pipeline_consumer_state, transform2mma_flag);
int transform2mma_pipeline_consumer_state_index = curr_transform2mma_pipeline_consumer_state.index(); //read_stage
int mma2accum_pipeline_producer_state_index = curr_mma2accum_pipeline_producer_state.index(); //write_stage
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
auto tCtC = accumulators(_,_,_,mma2accum_pipeline_producer_state_index);
auto tCrA0 = tCrA(_,_,_,transform2mma_pipeline_consumer_state_index);
auto tCrB0 = tCrB(_,_,_,transform2mma_pipeline_consumer_state_index);
cute::gemm(tiled_mma, tCrA0(_,_,k_block), tCrB0(_,_,k_block), tCtC); // A[0]*B[0]
tiled_mma.accumulate_ = UMMA::ScaleOut::One;
}
transform2mma_pipeline.consumer_release(curr_transform2mma_pipeline_consumer_state);
skip_wait = (k_tile_count <= 1);
transform2mma_flag = transform2mma_pipeline.consumer_try_wait(next_transform2mma_pipeline_consumer_state, skip_wait);
curr_transform2mma_pipeline_consumer_state = next_transform2mma_pipeline_consumer_state;
++next_transform2mma_pipeline_consumer_state;
}
mma2accum_pipeline.producer_commit(curr_mma2accum_pipeline_producer_state);
return cute::make_tuple(curr_transform2mma_pipeline_consumer_state, mma2accum_pipeline_producer_state);
}
template<class FrgEngine, class FrgLayout>
CUTLASS_DEVICE auto
mma_init(cute::Tensor<FrgEngine, FrgLayout> const& accumulators, TensorStorage& shared_storage) const {
TiledMma tiled_mma;
auto get_tCrA = [&] () constexpr {
if constexpr (cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value) {
Tensor sACompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), SmemLayoutACompute{});
return tiled_mma.make_fragment_A(sACompute);
}
else {
auto tCrA = tiled_mma.make_fragment_A(shape(SmemLayoutACompute{}));
tCrA.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators);
return tCrA;
}
};
Tensor tCrA = get_tCrA();
Tensor sB = make_tensor(make_smem_ptr(shared_storage.input.smem_B.begin()), SmemLayoutB{});
Tensor tCrB = tiled_mma.make_fragment_B(sB);
return cute::make_tuple(tCrA, tCrB);
}
template<class FrgEngine, class FrgLayout, class TmemCopyAtom, class EpilogueTile>
CUTLASS_DEVICE auto
accum_init(cute::Tensor<FrgEngine, FrgLayout> const& accumulators, TmemCopyAtom tmem_cp_atom, EpilogueTile epilogue_tile) {
return accumulators;
}
private:
template <class ProblemShape_MNKL>
CUTLASS_DEVICE
constexpr auto
tile_input_tensors(Params const& params, ProblemShape_MNKL const& problem_shape_MNKL) const {
using X = cute::Underscore;
// Separate out problem shape for convenience
auto [M,N,K,L] = problem_shape_MNKL;
// Represent the full tensors -- get these from TMA
Tensor mA_mkl = observed_tma_load_a_->get_tma_tensor(make_shape(M,K,L));
Tensor mB_nkl = observed_tma_load_b_->get_tma_tensor(make_shape(N,K,L));
// Tile the tensors and defer the slice
Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{});
Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{});
return cute::make_tuple(gA_mkl, gB_nkl);
}
typename Params::TMA_A const* observed_tma_load_a_ = nullptr;
typename Params::TMA_B const* observed_tma_load_b_ = nullptr;
ClusterShape cluster_shape_;
uint32_t block_rank_in_cluster_;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp
+12 -9
View File
@@ -353,13 +353,12 @@ struct CollectiveMma<
(void) workspace;
// These tensor shapes (only applicable for grouped gemm) and pointers are only used to create tensormap/tma desc.
// These will be replaced with correct values before the initial tma load.
auto init_shape = repeat_like(typename ProblemShape::UnderlyingProblemShape{}, int32_t(1));
constexpr int tma_alignment_bits = 128;
auto init_M = tma_alignment_bits;
auto init_N = tma_alignment_bits;
auto init_K = tma_alignment_bits;
auto init_M = int32_t(size<0>(TileShape{}));
auto init_N = int32_t(size<1>(TileShape{}));
auto init_K = int32_t(size<2>(TileShape{}));
auto init_L = 1;
// Batches/Groups are managed by using appropriate pointers to input matrices
const uint32_t init_L = 1;
TmaInternalElementA const* ptr_A_first_batch = nullptr;
TmaInternalElementB const* ptr_B_first_batch = nullptr;
ElementSF const* ptr_SFA_first_batch = nullptr;
@@ -1058,11 +1057,11 @@ struct CollectiveMma<
Tensor tensor_sfb = make_tensor(ptr_SF, mainloop_params.layout_SFB[next_group]);
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_a, tensor_a,
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_a, tensor_a,
prob_shape_A, prob_stride_A);
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_sfa, tensor_sfa,
prob_shape_SFA, prob_stride_SFA);
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_b, tensor_b,
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_b, tensor_b,
prob_shape_B, prob_stride_B);
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_sfb, tensor_sfb,
prob_shape_SFB, prob_stride_SFB);
@@ -1091,7 +1090,7 @@ struct CollectiveMma<
prob_stride_B);
cute::tma_descriptor_replace_dims_strides_in_shared_mem(shared_tensormaps.smem_tensormap_SFB,
prob_shape_SFB,
prob_stride_SFB);
prob_stride_SFB);
}
// The entire warp must call this function collectively (that is, the instructions are aligned)
@@ -1122,6 +1121,10 @@ struct CollectiveMma<
tensormaps_cp_fence_release (
TensorMapStorage& shared_tensormaps,
cute::tuple<TensorMapA, TensorMapB, TensorMapSFA, TensorMapSFB> const& input_tensormaps) {
if (cute::elect_one_sync()) {
cute::tma_desc_commit_group();
cute::tma_desc_wait_group();
}
// Entire warp must do this (i.e. it's aligned)
tma_descriptor_cp_fence_release(get<0>(input_tensormaps), shared_tensormaps.smem_tensormap_A);
tma_descriptor_cp_fence_release(get<1>(input_tensormaps), shared_tensormaps.smem_tensormap_B);
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input.hpp
+4
View File
@@ -1330,6 +1330,10 @@ public:
tensormaps_cp_fence_release (
TensorMapStorage& shared_tensormaps,
cute::tuple<TMs...> const& input_tensormaps) {
if (cute::elect_one_sync()) {
cute::tma_desc_commit_group();
cute::tma_desc_wait_group();
}
// Entire warp must do this (i.e. it's aligned)
tma_descriptor_cp_fence_release(get<0>(input_tensormaps), shared_tensormaps.smem_tensormap_A);
tma_descriptor_cp_fence_release(get<1>(input_tensormaps), shared_tensormaps.smem_tensormap_B);
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized.hpp
+5 -2
View File
@@ -529,10 +529,9 @@ struct CollectiveMma<
// Prologue GMMAs
int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
assert(k_tile_count >= 1);
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
warpgroup_fence_operand(accum);
{
if (k_tile_count > 0) {
// WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
pipeline.consumer_wait(smem_pipe_read, barrier_token);
@@ -739,6 +738,10 @@ struct CollectiveMma<
tensormaps_cp_fence_release (
TensorMapStorage& shared_tensormaps,
cute::tuple<TensorMapA, TensorMapB> const& input_tensormaps) {
if (cute::elect_one_sync()) {
cute::tma_desc_commit_group();
cute::tma_desc_wait_group();
}
// Entire warp must do this (i.e. it's aligned)
tma_descriptor_cp_fence_release(get<0>(input_tensormaps), shared_tensormaps.smem_tensormap_A);
tma_descriptor_cp_fence_release(get<1>(input_tensormaps), shared_tensormaps.smem_tensormap_B);
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8.hpp
+4
View File
@@ -747,6 +747,10 @@ struct CollectiveMma<
tensormaps_cp_fence_release (
TensorMapStorage& shared_tensormaps,
cute::tuple<TensorMapA, TensorMapB> const& input_tensormaps) {
if (cute::elect_one_sync()) {
cute::tma_desc_commit_group();
cute::tma_desc_wait_group();
}
// Entire warp must do this (i.e. it's aligned)
tma_descriptor_cp_fence_release(get<0>(input_tensormaps), shared_tensormaps.smem_tensormap_A);
tma_descriptor_cp_fence_release(get<1>(input_tensormaps), shared_tensormaps.smem_tensormap_B);
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+135 -142
View File
@@ -135,6 +135,10 @@ struct CollectiveMma<
static constexpr int ScalePromotionInterval = ScaleGranularityK / size<2>(typename TiledMma::AtomShape_MNK{});
static_assert(ScalePromotionInterval % 4 == 0, "ScalePromotionInterval must be a multiple of 4.");
static_assert(ScalePromotionInterval >= size<2>(TileShape{}) / tile_size<2>(TiledMma{}),
"ScalePromotionInterval must be greater than or equal to the number of stages of the MMA atom.");
static_assert(ScalePromotionInterval % (size<2>(TileShape{}) / tile_size<2>(TiledMma{})) == 0,
"ScalePromotionInterval must be a multiple of the number of stages of the MMA atom.");
static constexpr int ScaleMsPerTile = size<0>(TileShape{}) / ScaleGranularityM;
static constexpr int ScaleNsPerTile = size<1>(TileShape{}) / ScaleGranularityN;
@@ -811,31 +815,37 @@ struct CollectiveMma<
Tensor tCrSFB = make_tensor_like<ElementBlockScale>(tCsSFB(_, _, _, _0{})); // (MMA,MMA_M,MMA_N)
// Prologue GMMAs
int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
assert(k_tile_count >= 1);
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
// WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
pipeline.consumer_wait(smem_pipe_read, barrier_token);
// fence_operand();
GmmaFP8Accumulation accumulation(accum, ScalePromotionInterval, size<2>(tCrA));
{
// WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
pipeline.consumer_wait(smem_pipe_read, barrier_token);
if constexpr (ScalePromotionInterval != 4) {
if (accumulation.prepare_if_needed()) {
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
}
else {
// Always zero out the accumulator for finest granularity
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
warpgroup_fence_operand(accumulation());
{
int read_stage = smem_pipe_read.index();
// Load per block scale values from shared memory to registers
copy(tCsSFA(_,_,_,make_coord(_0{},read_stage)), tCrSFA);
copy(tCsSFB(_,_,_,make_coord(_0{},read_stage)), tCrSFB);
warpgroup_fence_operand(accumulation());
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
}
warpgroup_commit_batch();
warpgroup_fence_operand(accumulation());
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{});
}
@@ -854,16 +864,9 @@ struct CollectiveMma<
}
}
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
}
warpgroup_commit_batch();
warpgroup_wait<0>();
++smem_pipe_read;
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
// Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB`
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
@@ -879,90 +882,16 @@ struct CollectiveMma<
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFA, tCrSFB);
}
++smem_pipe_read;
}
warpgroup_fence_operand(accumulation());
CUTLASS_PRAGMA_UNROLL
for (int k_tile_prologue = prologue_mma_count - 1; k_tile_prologue > 0; --k_tile_prologue)
{
// WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
pipeline.consumer_wait(smem_pipe_read, barrier_token);
if constexpr (ScalePromotionInterval != 4) {
if (accumulation.prepare_if_needed()) {
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
}
else {
// Always zero out the accumulator for finest granularity
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
int read_stage = smem_pipe_read.index();
// Load per block scale values from shared memory to registers
copy(tCsSFA(_,_,_,make_coord(_0{},read_stage)), tCrSFA);
copy(tCsSFB(_,_,_,make_coord(_0{},read_stage)), tCrSFB);
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{});
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_b = tCrSFB(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) {
filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b;
}
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
ElementBlockScale scale_a = tCrSFA(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) {
filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a;
}
}
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
}
warpgroup_commit_batch();
// Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB`
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA(_0{});
scale_if_needed(accumulation, scale_ab);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
scale_if_needed(accumulation, tCrSFA);
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFB);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFA, tCrSFB);
}
++smem_pipe_read;
}
warpgroup_fence_operand(accumulation());
// Mainloop GMMAs
k_tile_count -= prologue_mma_count;
k_tile_count--;
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 0; --k_tile_count)
for ( ; k_tile_count > 1; --k_tile_count)
{
// WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
pipeline.consumer_wait(smem_pipe_read, barrier_token);
//
@@ -970,29 +899,10 @@ struct CollectiveMma<
//
int read_stage = smem_pipe_read.index();
// fence_operand();
// Load per block scale values from shared memory to registers (at most twice per block along M and/or N)
copy(tCsSFA(_,_,_,make_coord(_0{},read_stage)), tCrSFA);
copy(tCsSFB(_,_,_,make_coord(_0{},read_stage)), tCrSFB);
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{});
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_b = tCrSFB(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) {
filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b;
}
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
ElementBlockScale scale_a = tCrSFA(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) {
filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a;
}
}
copy(tCsSFA(_,_,_,make_coord(_0{}, read_stage)), tCrSFA);
copy(tCsSFB(_,_,_,make_coord(_0{}, read_stage)), tCrSFB);
if constexpr (ScalePromotionInterval != 4) {
if (accumulation.prepare_if_needed()) {
@@ -1005,7 +915,6 @@ struct CollectiveMma<
}
warpgroup_fence_operand(accumulation());
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
@@ -1017,9 +926,31 @@ struct CollectiveMma<
warpgroup_commit_batch();
/// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed
warpgroup_wait<K_PIPE_MMAS>();
warpgroup_fence_operand(accumulation());
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{});
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_b = tCrSFB(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) {
filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b;
}
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
ElementBlockScale scale_a = tCrSFA(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) {
filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a;
}
}
warpgroup_wait<0>();
pipeline.consumer_release(smem_pipe_release); // Unlock previous tile
++smem_pipe_read;
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
// Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB`
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA(_0{});
@@ -1035,13 +966,81 @@ struct CollectiveMma<
scale_if_needed(accumulation, tCrSFA, tCrSFB);
}
// UNLOCK smem_pipe_release, done _computing_ on it
pipeline.consumer_release(smem_pipe_release);
// Advance smem_pipe_read and smem_pipe_release
++smem_pipe_read;
++smem_pipe_release;
}
if (k_tile_count) {
pipeline.consumer_wait(smem_pipe_read, barrier_token);
//
// Compute on k_tile
//
int read_stage = smem_pipe_read.index();
// Load per block scale values from shared memory to registers (at most twice per block along M and/or N)
copy(tCsSFA(_,_,_,make_coord(_0{}, read_stage)), tCrSFA);
copy(tCsSFB(_,_,_,make_coord(_0{}, read_stage)), tCrSFB);
if constexpr (ScalePromotionInterval != 4) {
if (accumulation.prepare_if_needed()) {
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
}
else {
// Always zero out the accumulator for finest granularity
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
warpgroup_fence_operand(accumulation());
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
}
warpgroup_commit_batch();
/// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed
warpgroup_fence_operand(accumulation());
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{});
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_b = tCrSFB(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) {
filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b;
}
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
ElementBlockScale scale_a = tCrSFA(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) {
filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a;
}
}
warpgroup_wait<0>();
pipeline.consumer_release(smem_pipe_release); // Unlock previous tile
// Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB`
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA(_0{});
scale_if_needed(accumulation, scale_ab);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
scale_if_needed(accumulation, tCrSFA);
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFB);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFA, tCrSFB);
}
}
if constexpr (ScalePromotionInterval != 4) {
// residues only exists when granularity is not the finnest
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
@@ -1066,19 +1065,9 @@ struct CollectiveMma<
/// Perform a Consumer Epilogue to release all buffers
CUTLASS_DEVICE void
mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) {
// Prologue GMMAs
int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
k_tile_count -= prologue_mma_count;
smem_pipe_release.advance(k_tile_count);
// Wait on all GMMAs to complete
warpgroup_wait<0>();
for (int count = 0; count < prologue_mma_count; ++count) {
pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it
++smem_pipe_release;
}
// The pipeline is not released in the first iteration
smem_pipe_release.advance(k_tile_count - 1);
pipeline.consumer_release(smem_pipe_release);
}
//
@@ -1198,6 +1187,10 @@ struct CollectiveMma<
tensormaps_cp_fence_release (
TensorMapStorage& shared_tensormaps,
cute::tuple<TensorMapA, TensorMapB> const& input_tensormaps) {
if (cute::elect_one_sync()) {
cute::tma_desc_commit_group();
cute::tma_desc_wait_group();
}
// Entire warp must do this (i.e. it's aligned)
tma_descriptor_cp_fence_release(get<0>(input_tensormaps), shared_tensormaps.smem_tensormap_A);
tma_descriptor_cp_fence_release(get<1>(input_tensormaps), shared_tensormaps.smem_tensormap_B);
include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+5 -12
View File
@@ -128,6 +128,10 @@ struct CollectiveMma<
static constexpr int ScalePromotionInterval = ScaleGranularityK / size<2>(typename TiledMma::AtomShape_MNK{});
static_assert(ScalePromotionInterval % 4 == 0, "ScalePromotionInterval must be a multiple of 4.");
static_assert(ScalePromotionInterval >= size<2>(TileShape{}) / tile_size<2>(TiledMma{}),
"ScalePromotionInterval must be greater than or equal to the number of stages of the MMA atom.");
static_assert(ScalePromotionInterval % (size<2>(TileShape{}) / tile_size<2>(TiledMma{})) == 0,
"ScalePromotionInterval must be a multiple of the number of stages of the MMA atom.");
static constexpr int ScaleMsPerTile = size<0>(TileShape{}) / ScaleGranularityM;
static constexpr int ScaleNsPerTile = size<1>(TileShape{}) / ScaleGranularityN;
@@ -213,7 +217,6 @@ struct CollectiveMma<
StrideA dA;
ElementB const* ptr_B;
StrideB dB;
uint32_t mma_promotion_interval = 4;
ElementBlockScale const* ptr_SFA;
LayoutSFA layout_SFA;
ElementBlockScale const* ptr_SFB;
@@ -382,16 +385,6 @@ struct CollectiveMma<
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem size doesn't meet the minimum alignment requirements for using TMA to load scale B.\n");
}
/* MMA promotion interval should be a multiple of 4, since each mainloop iteration would issue 4 MMA instructions. */
constexpr int pipe_k = size<2>(TileShape{}) / tile_size<2>(TiledMma{});
if (args.mma_promotion_interval % 4 != 0 ||
args.mma_promotion_interval != ScalePromotionInterval ||
args.mma_promotion_interval % pipe_k != 0 ||
pipe_k > args.mma_promotion_interval) {
implementable = false;
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Argument mma_promotion_interval is invalid.\n");
}
// We expect full tiles in K
if (K % size<2>(TileShape{}) != 0) {
implementable = false;
@@ -1001,7 +994,7 @@ struct CollectiveMma<
// Advance smem_pipe_read and smem_pipe_release
++smem_pipe_release;
}
if (k_tile_count == 1) {
if (k_tile_count) {
pipeline.consumer_wait(smem_pipe_read, barrier_token);
//
include/cutlass/gemm/dispatch_policy.hpp
+2
View File
@@ -1016,6 +1016,8 @@ struct MainloopSm100ArrayTmaUmmaWarpSpecializedFastF32 {
// Mainloop schedule for array-based TMA
template<
int Stages_,
int SchedulerPipelineStageCount_,
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp
+5 -3
View File
@@ -824,9 +824,6 @@ public:
}
else if (is_participant.sched) {
if constexpr (IsSchedDynamicPersistent) {
cutlass::arch::wait_on_dependent_grids();
}
// Signal the epilogue warps to proceed once the prologue is complete
epilogue_throttle_barrier.arrive();
@@ -837,6 +834,8 @@ public:
// why this variable is needed.
bool requires_clc_query = true;
cutlass::arch::wait_on_dependent_grids();
do {
if (requires_clc_query) {
// Throttle CLC query to mitigate workload imbalance caused by skews among persistent workers.
@@ -872,6 +871,9 @@ public:
clc_pipeline.producer_tail(clc_pipe_producer_state);
}
else {
cutlass::arch::wait_on_dependent_grids();
do {
auto [next_work_tile_info, increment_pipe] = scheduler.advance_to_next_work(clc_pipeline, clc_pipe_producer_state);
work_tile_info = next_work_tile_info;
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp
+3 -4
View File
@@ -832,10 +832,6 @@ public:
// Register reconfiguration
arch::warpgroup_reg_dealloc<GenericRegisterRequirement>();
if constexpr (IsSchedDynamicPersistent) {
cutlass::arch::wait_on_dependent_grids();
}
// Signal the epilogue warps to proceed once the prologue is complete
epilogue_throttle_barrier.arrive();
@@ -845,6 +841,8 @@ public:
// See comment below where this variable is updated for a description of
// why this variable is needed.
bool requires_clc_query = true;
cutlass::arch::wait_on_dependent_grids();
do {
if (requires_clc_query) {
// Throttle CLC query to mitigate workload imbalance caused by skews among persistent workers.
@@ -883,6 +881,7 @@ public:
clc_pipeline.producer_tail(clc_pipeline_producer_state);
}
else {
cutlass::arch::wait_on_dependent_grids();
do {
auto [next_work_tile_info, increment_pipe] = scheduler.advance_to_next_work(clc_pipeline, clc_pipeline_producer_state);
work_tile_info = next_work_tile_info;
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_mma_transform.hpp
+3
View File
@@ -941,6 +941,8 @@ public:
// why this variable is needed.
bool requires_clc_query = true;
cutlass::arch::wait_on_dependent_grids();
do {
if (requires_clc_query) {
// Throttle CLC query to mitigate workload imbalance caused by skews among persistent workers.
@@ -976,6 +978,7 @@ public:
clc_pipeline.producer_tail(clc_pipe_producer_state);
}
else {
cutlass::arch::wait_on_dependent_grids();
do {
auto [next_work_tile_info, increment_pipe] = scheduler.advance_to_next_work(clc_pipeline, clc_pipe_producer_state);
work_tile_info = next_work_tile_info;
include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized.hpp
+2 -4
View File
@@ -177,7 +177,6 @@ public:
// Kernel level shared memory storage
struct SharedStorage {
// Barriers should be allocated in lower 8KB of SMEM for SM100
struct PipelineStorage : cute::aligned_struct<16, _1> {
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
@@ -649,15 +648,14 @@ public:
}
else if (is_participant.sched) {
if constexpr (IsSchedDynamicPersistent) {
cutlass::arch::wait_on_dependent_grids();
}
if constexpr (IsSchedDynamicPersistent) {
// Whether a new CLC query must be performed.
// See comment below where this variable is updated for a description of
// why this variable is needed.
bool requires_clc_query = true;
cutlass::arch::wait_on_dependent_grids();
do {
if (requires_clc_query) {
// Query next clcID and update producer state
include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp
+3 -4
View File
@@ -717,10 +717,6 @@ public:
// Register reconfiguration
arch::warpgroup_reg_dealloc<GenericRegisterRequirement>();
if constexpr (IsSchedDynamicPersistent) {
cutlass::arch::wait_on_dependent_grids();
}
// Signal the epilogue warps to proceed once the prologue is complete
epilogue_throttle_barrier.arrive();
@@ -729,6 +725,9 @@ public:
// See comment below where this variable is updated for a description of
// why this variable is needed.
bool requires_clc_query = true;
cutlass::arch::wait_on_dependent_grids();
do {
if (requires_clc_query) {
// Throttle CLC query to mitigate workload imbalance caused by skews among persistent workers.
include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_mma_transform.hpp
+2 -4
View File
@@ -769,10 +769,6 @@ public:
// Register reconfiguration
arch::warpgroup_reg_dealloc<GenericRegisterRequirement>();
if constexpr (IsSchedDynamicPersistent) {
cutlass::arch::wait_on_dependent_grids();
}
// Signal the epilogue warps to proceed once the prologue is complete
epilogue_throttle_barrier.arrive();
@@ -783,6 +779,8 @@ public:
// why this variable is needed.
bool requires_clc_query = true;
cutlass::arch::wait_on_dependent_grids();
do {
if (requires_clc_query) {
// Throttle CLC query to mitigate workload imbalance caused by skews among persistent workers.
include/cutlass/gemm/kernel/sm100_sparse_gemm_tma_warpspecialized.hpp
+2 -3
View File
@@ -687,15 +687,14 @@ public:
}
else if (is_participant.sched) {
if constexpr (IsSchedDynamicPersistent) {
cutlass::arch::wait_on_dependent_grids();
}
if constexpr (IsSchedDynamicPersistent) {
// Whether a new CLC query must be performed.
// See comment below where this variable is updated for a description of
// why this variable is needed.
bool requires_clc_query = true;
cutlass::arch::wait_on_dependent_grids();
do {
if (requires_clc_query) {
// Query next clcID and update producer state
include/cutlass/gemm/kernel/sm120_gemm_tma_warpspecialized_cooperative_asymmetric_dma.hpp
+3 -4
View File
@@ -551,13 +551,12 @@ public:
if (producer_warp_role == ProducerWarpRole::Warp1) {
work_tile_info = scheduler.initial_work_tile_info(ClusterShape{});
if constexpr (IsSchedDynamicPersistent) {
cutlass::arch::wait_on_dependent_grids();
}
if constexpr (IsSchedDynamicPersistent) {
bool requires_clc_query = true;
TileSchedulerPipelineState scheduler_pipe_producer_state = cutlass::make_producer_start_state<TileSchedulerPipeline>();
cutlass::arch::wait_on_dependent_grids();
while (work_tile_info.is_valid()) {
if (requires_clc_query) {
// Throttle CLC query to mitigate workload imbalance caused by skews among persistent workers.
include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp
+13 -21
View File
@@ -153,12 +153,12 @@ public:
cute::conditional_t<cute::is_same_v<SchedulerTag, void>, void, ProblemShape> // Use void for default scheduler.
>::Scheduler;
static constexpr auto TileSchedulerStages = 8;
using TileSchedulerArguments = typename TileScheduler::Arguments;
using TileSchedulerParams = typename TileScheduler::Params;
using TileSchedulerResponse = typename TileSchedulerResponseGetter<TileScheduler>::Type;
static constexpr auto TileSchedulerStages = 8;
static constexpr uint32_t NumLoadWarpGroups = 1;
static constexpr uint32_t NumMmaThreads = size(TiledMma{});
static constexpr uint32_t NumMmaWarpGroups = NumMmaThreads / NumThreadsPerWarpGroup;
@@ -378,7 +378,6 @@ public:
if (status != Status::kSuccess) {
return status;
}
return status;
}
@@ -461,6 +460,7 @@ public:
return TileScheduler{params.scheduler};
}
} ();
// In a warp specialized kernel, collectives expose data movement and compute operations separately
CollectiveMainloop collective_mainloop;
CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
@@ -484,7 +484,7 @@ public:
typename TileSchedulerPipeline::Params tile_scheduler_pipeline_params;
if constexpr (cute::is_same_v<SchedulerTag, GroupScheduler>) {
if (warp_group_role == WarpGroupRole::Producer
&& producer_warp_role == ProducerWarpRole::Scheduler) {
&& producer_warp_role == ProducerWarpRole::Scheduler) {
tile_scheduler_pipeline_params.role = TileSchedulerPipeline::ThreadCategory::Producer;
}
else {
@@ -499,7 +499,6 @@ public:
tile_scheduler_pipeline_params.producer_arv_count = 1;
}
TileSchedulerPipeline tile_scheduler_pipeline(shared_storage.pipelines.scheduler, tile_scheduler_pipeline_params);
// Mainloop Load pipeline
using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
typename MainloopPipeline::Params mainloop_pipeline_params;
@@ -683,9 +682,8 @@ public:
block_rank_in_cluster,
shared_storage.tensors.mainloop
);
// Update starting pipeline state for the next tile
// Wait for the last TMA stage to complete loading, before issuing tensormap updates
mainloop_pipe_producer_state.advance(work_k_tile_count - 1);
// Pipeline state is only advanced if there are K tiles to compute
mainloop_pipe_producer_state.advance(work_k_tile_count);
// Signal for the epilogue load warp to begin
if (do_load_order_arrive) {
@@ -706,11 +704,6 @@ public:
if constexpr (IsGroupedGemmKernel) {
problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(curr_batch), 1);
}
// Purpose of this pipeline state is to make sure TMA loads have finished before doing descriptor updates
// Since this state is waiting for loads to finish, it must start in the inverted phase.
typename CollectiveMainloop::PipelineState mainloop_pipe_tma_consumer_state =
{mainloop_pipe_producer_state.index(), !mainloop_pipe_producer_state.phase(), mainloop_pipe_producer_state.count()};
mainloop_pipeline.consumer_wait(mainloop_pipe_tma_consumer_state);
collective_mainloop.tensormaps_perform_update(
shared_storage.tensormaps.mainloop,
params.mainloop,
@@ -723,8 +716,6 @@ public:
// Entire warp must do this (i.e. it's aligned)
collective_mainloop.tensormaps_cp_fence_release(shared_storage.tensormaps.mainloop, input_tensormaps);
}
// Advance the producer state for the last remaining stage that was being waited for above
mainloop_pipe_producer_state.advance(1);
} while (work_tile_info.is_valid()); // Scheduler work fetch loop
// Make sure all Consumer Warp Groups have been waited upon
@@ -771,8 +762,8 @@ public:
block_rank_in_cluster,
shared_storage.tensors.mainloop
);
// Update starting pipeline state for the next tile
// Wait for the last TMA stage to complete loading, before issuing tensormap updates
mainloop_pipe_producer_state.advance(work_k_tile_count);
// Get next work tile
@@ -790,8 +781,8 @@ public:
}
}
} while (work_tile_info.is_valid()); // Scheduler work fetch loop
}
}
} // End of auxiliary load needed check
} // Mainloop Auxiliary Load Producer Warp End
// Epilogue Producer Warp
else if (producer_warp_role == ProducerWarpRole::Epilogue && collective_epilogue.is_producer_load_needed()) {
int32_t const sm_idx = blockIdx.x + (blockIdx.y * gridDim.x);
@@ -854,6 +845,7 @@ public:
wait
);
}
work_tile_info = next_work_tile_info;
if (increment_pipe) {
++tile_scheduler_pipe_consumer_state;
@@ -917,8 +909,8 @@ public:
// Converge before issuing tensormap fence release since fence is aligned
__syncwarp();
collective_epilogue.template tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue,
epi_store_tensormap,
collective_epilogue.template tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue,
epi_store_tensormap,
consumer_warp_group_idx);
}
@@ -1021,7 +1013,7 @@ public:
// Converge before issuing tensormap fence release since fence is aligned
__syncwarp();
collective_epilogue.template tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue,
collective_epilogue.template tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue,
epi_store_tensormap,
consumer_warp_group_idx);
}
include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp
+6 -13
View File
@@ -69,7 +69,6 @@ class GemmUniversal<
cute::enable_if_t<cute::is_base_of_v<KernelPtrArrayTmaWarpSpecializedPingpong, typename CollectiveMainloop_::DispatchPolicy::Schedule>>
>
{
// Get the type of the scheduler response.
template<typename TileScheduler, typename = void>
struct TileSchedulerResponseGetter {
@@ -145,7 +144,6 @@ public:
TileScheduler_
>;
using TileScheduler = typename detail::TileSchedulerSelector<
SchedulerTag,
ArchTag,
@@ -646,6 +644,8 @@ public:
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
// the work tile infos to the downstream pipelines.
if constexpr (cute::is_same_v<SchedulerTag, GroupScheduler>) {
do {
auto [next_work_tile_info, increment_pipe] = scheduler.advance_to_next_work(tile_scheduler_pipeline, tile_scheduler_pipe_producer_state);
@@ -684,7 +684,8 @@ public:
bool did_batch_change = true;
do {
if (!TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info, tile_scheduler_pipeline, tile_scheduler_pipe_consumer_state);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
work_tile_info, tile_scheduler_pipeline, tile_scheduler_pipe_consumer_state);
work_tile_info = next_work_tile_info;
if (increment_pipe) {
++tile_scheduler_pipe_consumer_state;
@@ -719,9 +720,8 @@ public:
block_rank_in_cluster,
shared_storage.tensors.mainloop
);
// Update starting pipeline state for the next tile
// Wait for the last TMA stage to complete loading, before issuing tensormap updates
mainloop_pipe_producer_state.advance(work_k_tile_count - 1);
// Pipeline state is only advanced if there are K tiles to compute
mainloop_pipe_producer_state.advance(work_k_tile_count);
// Signal for the epilogue load warp to begin
if (do_load_order_arrive) {
@@ -742,11 +742,6 @@ public:
if constexpr (IsGroupedGemmKernel) {
problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(curr_batch), 1);
}
// Purpose of this pipeline state is to make sure TMA loads have finished before doing descriptor updates
// Since this state is waiting for loads to finish, it must start in the inverted phase.
typename CollectiveMainloop::PipelineState mainloop_pipe_tma_consumer_state =
{mainloop_pipe_producer_state.index(), !mainloop_pipe_producer_state.phase(), mainloop_pipe_producer_state.count()};
mainloop_pipeline.consumer_wait(mainloop_pipe_tma_consumer_state);
collective_mainloop.tensormaps_perform_update(
shared_storage.tensormaps.mainloop,
params.mainloop,
@@ -759,8 +754,6 @@ public:
// Entire warp must do this (i.e. it's aligned)
collective_mainloop.tensormaps_cp_fence_release(shared_storage.tensormaps.mainloop, input_tensormaps);
}
// Advance the producer state for the last remaining stage that was being waited for above
mainloop_pipe_producer_state.advance(1);
} while (work_tile_info.is_valid()); // Scheduler work fetch loop
// Make sure all Consumer Warp Groups have been waited upon
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp
+2
View File
@@ -546,6 +546,8 @@ public:
if constexpr (IsSchedDynamicPersistent) {
bool requires_clc_query = true;
TileSchedulerPipelineState scheduler_pipe_producer_state = cutlass::make_producer_start_state<TileSchedulerPipeline>();
cutlass::arch::wait_on_dependent_grids();
while (work_tile_info.is_valid()) {
if (requires_clc_query) {
include/cutlass/pipeline/sm100_pipeline.hpp
+3 -3
View File
@@ -334,9 +334,9 @@ public:
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
static_assert(IsDynamicCluster or ((cute::size<0>(cluster_shape) % cute::size<0>(atom_thr_shape) == 0) &&
(cute::size<1>(cluster_shape) % cute::size<1>(atom_thr_shape) == 0)));
uint32_t const multicast_consumer_arrival_count = (cute::size<0>(cluster_shape) / cute::size<0>(atom_thr_shape)) +
(cute::size<1>(cluster_shape) / cute::size<1>(atom_thr_shape)) - 1;
uint32_t const num_consumer_per_cluster = params.num_consumers / NumThreadsPerWarpGroup;
uint32_t const multicast_consumer_arrival_count = ((cute::size<0>(cluster_shape) / cute::size<0>(atom_thr_shape)) +
(cute::size<1>(cluster_shape) / cute::size<1>(atom_thr_shape)) - 1) * num_consumer_per_cluster;
cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(storage.full_barrier_), decltype(storage.empty_barrier_), Stages>(
storage.full_barrier_, storage.empty_barrier_, producer_arv_cnt, multicast_consumer_arrival_count);
}
include/cutlass/version.h
+3 -3
View File
@@ -34,9 +34,9 @@
#include <cstdint>
#include <string>
#define CUTLASS_MAJOR 3
#define CUTLASS_MINOR 9
#define CUTLASS_PATCH 2
#define CUTLASS_MAJOR 4
#define CUTLASS_MINOR 0
#define CUTLASS_PATCH 0
#ifdef CUTLASS_VERSIONS_GENERATED
#include "cutlass/version_extended.h"