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

v4.3 update. (#2709)

Browse Source 
* v4.3 update.

* Update the cute_dsl_api changelog's doc link

* Update version to 4.3.0

* Update the example link

* Update doc to encourage user to install DSL from requirements.txt

---------

Co-authored-by: Larry Wu <larwu@nvidia.com>
This commit is contained in:
Junkai-Wu
2025-10-21 14:26:30 -04:00
committed by GitHub
co-authored by Larry Wu
parent e6e2cc29f5
commit b1d6e2c9b3
244 changed files with 59272 additions and 10455 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl
+8 -3
View File
@@ -106,6 +106,7 @@ sm100_compute_stage_count_or_override_fast_fp32(StageCountAutoCarveout<carveout_
// FastFP (9xBF16) MMA kernels builder
template <
class ArchTag,
class GmemLayoutATag,
int AlignmentA,
class GmemLayoutBTag,
@@ -117,7 +118,7 @@ template <
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassTensorOp,
float, // ElementA
GmemLayoutATag, // LayoutA
@@ -131,6 +132,8 @@ struct CollectiveBuilder<
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
) &&
(not cute::is_tuple<GmemLayoutATag>::value && not cute::is_tuple<GmemLayoutBTag>::value) &&
(cute::is_base_of_v<KernelScheduleSm100FastFP32Gemm, BuilderScheduleTag>) &&
((sizeof(float) * AlignmentA) % detail::tma_alignment_bytes == 0) &&
@@ -226,9 +229,11 @@ struct CollectiveBuilder<
TensorMapStorage);
// Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations
static constexpr int Sm100ReducedSmemCapacityBytes = detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr auto stage_info = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_fast_fp32<
Sm100ReducedSmemCapacityBytes, CtaTileShape_MNK, TiledMma, BuilderScheduleTag, UmmaMajorA>(StageCountType{});
ReducedSmemCapacityBytes, CtaTileShape_MNK, TiledMma, BuilderScheduleTag, UmmaMajorA>(StageCountType{});
static constexpr int Load2TransformPipelineStageCount = get<0>(stage_info);
static constexpr int Transform2MmaPipelineStageCount = get<1>(stage_info);
include/cutlass/gemm/collective/builders/sm100_blockscaled_sparse_umma_builder.inl
+12 -2
View File
@@ -44,6 +44,7 @@ namespace detail {
// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count.
template <
int CapacityBytes,
class ElementAMma,
class ElementB,
class ElementEMma,
@@ -62,6 +63,7 @@ sm100_compute_stage_count_or_override_blockscaled_sparse(StageCount<stages> stag
// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count.
template <
int CapacityBytes,
class ElementAMma,
class ElementB,
class ElementEMma,
@@ -110,7 +112,7 @@ sm100_compute_stage_count_or_override_blockscaled_sparse(StageCountAutoCarveout<
constexpr auto EpilogueSharedStorage = carveout_bytes;
constexpr auto Stages = (cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout - EpilogueSharedStorage) /
constexpr auto Stages = (CapacityBytes - KernelSmemCarveout - EpilogueSharedStorage) /
(MainloopTensorStorage_per_Stage + MainloopPipelineStorage_per_Stage_aligned);
return Stages;
@@ -121,6 +123,7 @@ sm100_compute_stage_count_or_override_blockscaled_sparse(StageCountAutoCarveout<
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ArchTag,
class ElementPairA,
class GmemLayoutATag,
int AlignmentA,
@@ -134,7 +137,7 @@ template <
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassBlockScaledSparseTensorOp,
ElementPairA,
GmemLayoutATag,
@@ -148,6 +151,8 @@ struct CollectiveBuilder<
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
) &&
// Blockscaled Sparse Gemm
cute::is_base_of_v<KernelScheduleBlockScaledSparseGemmSm100, BuilderScheduleTag>
&&
@@ -272,7 +277,12 @@ struct CollectiveBuilder<
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
// Calculate SMEM capacity based on ArchTag
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes;
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_blockscaled_sparse<
ReducedSmemCapacityBytes,
ElementAMma_SmemAllocType,
ElementBMma_SmemAllocType,
ElementEMma,
include/cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl
+8 -3
View File
@@ -92,6 +92,7 @@ sm100_compute_stage_count_or_override_blockscaled(StageCountAutoCarveout<carveou
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ArchTag,
class ElementPairA,
class GmemLayoutATag,
int AlignmentA,
@@ -105,7 +106,7 @@ template <
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassBlockScaledTensorOp,
ElementPairA,
GmemLayoutATag,
@@ -119,6 +120,8 @@ struct CollectiveBuilder<
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
) &&
// Blockscaled Gemm
(not cute::is_same_v<KernelMixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100, BuilderScheduleTag>) &&
(cute::is_base_of_v<KernelScheduleBlockScaledGemmSm100, BuilderScheduleTag> ||
@@ -250,12 +253,14 @@ struct CollectiveBuilder<
4 // 4 Tensor maps for A, SFA, B and SFB
>::KernelSmemCarveout;
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int Sm100ReducedSmemCapacityBytes = cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_blockscaled<
Sm100ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, SmemLayoutAtomSFA, SmemLayoutAtomSFB>(StageCountType{});
ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, SmemLayoutAtomSFA, SmemLayoutAtomSFB>(StageCountType{});
static_assert(PipelineStages > 0, "Smem usage is too high. Can't create any SMEM buffers for A, B, SFA, and SFB.");
using DispatchPolicy =
include/cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl
+8 -3
View File
@@ -237,6 +237,7 @@ sm100_make_trivial_tiled_mma_blockwise() {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ArchTag,
class ElementA,
class GmemLayoutATagPair,
int AlignmentA,
@@ -250,7 +251,7 @@ template <
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassTensorOp,
ElementA,
GmemLayoutATagPair,
@@ -264,6 +265,8 @@ struct CollectiveBuilder<
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
) &&
not cute::is_tuple_v<ElementA> && not cute::is_tuple_v<ElementB> &&
not cute::is_complex_v<ElementA> && not cute::is_complex_v<ElementB> &&
cute::is_tuple_v<GmemLayoutATagPair> && cute::is_tuple_v<GmemLayoutBTagPair> &&
@@ -369,7 +372,9 @@ struct CollectiveBuilder<
IsArrayOfPointersGemm
>::KernelSmemCarveout;
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int Sm100ReducedSmemCapacityBytes = cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
using MainloopABPipelineStorage = typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage;
@@ -399,7 +404,7 @@ struct CollectiveBuilder<
using ScaleTileShape = cute::Shape<BlockTileScale_M, BlockTileScale_N, BlockTileScale_K>;
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_blockwise<
Sm100ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType,
ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType,
ElementAccumulator, ScaleTileShape, SmemTileShape, MainloopABPipelineStorage,
MainloopSFPipelineStorage>(StageCountType{});
static_assert(PipelineStages > 0, "Smem usage is too high. Can't create any SMEM buffers for A, B, and scales.");
include/cutlass/gemm/collective/builders/sm100_cpasync_umma_builder.inl
+13 -7
View File
@@ -39,6 +39,7 @@ namespace cutlass::gemm::collective {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ArchTag,
class ElementA,
class GmemLayoutATag,
int AlignmentA,
@@ -52,7 +53,7 @@ template <
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassTensorOp,
ElementA,
GmemLayoutATag,
@@ -65,10 +66,13 @@ struct CollectiveBuilder<
ClusterShape_MNK, // Static cluster shape or dynamic (int, int, _1)
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<cute::is_same_v<KernelWarpSpecialized1SmSm100, BuilderScheduleTag> ||
(cute::is_same_v<KernelScheduleAuto, BuilderScheduleTag> &&
(((sizeof(ElementA) * AlignmentA) % cutlass::gemm::collective::detail::tma_alignment_bytes != 0) ||
((sizeof(ElementB) * AlignmentB) % cutlass::gemm::collective::detail::tma_alignment_bytes != 0)))>
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
) &&
(cute::is_same_v<KernelWarpSpecialized1SmSm100, BuilderScheduleTag> ||
(cute::is_same_v<KernelScheduleAuto, BuilderScheduleTag> &&
(((sizeof(ElementA) * AlignmentA) % cutlass::gemm::collective::detail::tma_alignment_bytes != 0) ||
((sizeof(ElementB) * AlignmentB) % cutlass::gemm::collective::detail::tma_alignment_bytes != 0))))>
>
{
static_assert(cute::is_static_v<TileShape_MNK>, "TileShape has to be static");
@@ -137,13 +141,15 @@ struct CollectiveBuilder<
CLCPipelineStorage +
CLCResponseStorage);
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int Sm100ReducedSmemCapacityBytes = cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
using MainloopPipelineStorage = typename cutlass::PipelineUmmaConsumerAsync<1>::SharedStorage;
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override<
Sm100ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, MainloopPipelineStorage>(StageCountType{});
ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, MainloopPipelineStorage>(StageCountType{});
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
cutlass::gemm::MainloopSm100UmmaCpAsyncWarpSpecialized<
include/cutlass/gemm/collective/builders/sm100_mixed_input_umma_builder.inl
+7 -4
View File
@@ -152,6 +152,7 @@ constexpr int get_ScaleGranularityK() {
// Mixed Input MMA kernels builder
template <
class ArchTag,
class ElementAOptionalTuple,
class GmemLayoutATagTuple,
int AlignmentA,
@@ -165,7 +166,7 @@ template <
class KernelScheduleType
>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassTensorOp,
ElementAOptionalTuple, // ElementA
GmemLayoutATagTuple, // LayoutA
@@ -179,6 +180,8 @@ struct CollectiveBuilder<
StageCountType,
KernelScheduleType,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
) &&
(cute::is_base_of_v<KernelScheduleSm100MixedInputGemm, KernelScheduleType>) &&
((sizeof(float) * AlignmentA) % detail::tma_alignment_bytes == 0) &&
((sizeof(float) * AlignmentB) % detail::tma_alignment_bytes == 0)>>
@@ -304,12 +307,12 @@ struct CollectiveBuilder<
TensorMapStorage);
// Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations
static constexpr int Sm100ReducedSmemCapacityBytes = detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ScaleGranularityK = get_ScaleGranularityK<LayoutScale>();
static constexpr auto stage_info = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_mixed_input<
Sm100ReducedSmemCapacityBytes, TmaElementA, ElementAMma, ElementScale, ElementZero, ElementB, CtaTileShape_MNK, TiledMma, KernelScheduleType, UmmaMajorA, ScaleGranularityK>(StageCountType{});
ReducedSmemCapacityBytes, TmaElementA, ElementAMma, ElementScale, ElementZero, ElementB, CtaTileShape_MNK, TiledMma, KernelScheduleType, UmmaMajorA, ScaleGranularityK>(StageCountType{});
static constexpr int Load2TransformPipelineStageCount = get<0>(stage_info);
static constexpr int Transform2MmaPipelineStageCount = get<1>(stage_info);
include/cutlass/gemm/collective/builders/sm100_pipeline_carveout.inl
+13 -4
View File
@@ -46,11 +46,13 @@ struct Sm100DenseGemmTmaUmmaCarveout {
// AccumulatorPipeline = PipelineUmmaAsync
static constexpr auto AccumulatorPipelineStorage = sizeof(typename cutlass::PipelineUmmaAsync<AccumulatorPipelineStageCount>::SharedStorage);
// CLCPipeline = PipelineCLCFetchAsync
static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape_MNK>::SharedStorage);
// For pointer-array and grouped GEMM, we have two CLC responses, one for TMA updater, one for the TMA/MMA/Epilogue warps.
static constexpr int NumCLCResponses = (IsArrayOfPointersGemm ? 2 : 1);
static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape_MNK>::SharedStorage) * NumCLCResponses;
// LoadOrderBarrier = OrderedSequenceBarrier<1,2>
static constexpr auto LoadOrderBarrierStorage = sizeof(typename cutlass::OrderedSequenceBarrier<1,2>::SharedStorage);
// CLC (scheduler) response
static constexpr auto CLCResponseStorage = SchedulerPipelineStageCount * detail::CLCResponseSize;
static constexpr auto CLCResponseStorage = SchedulerPipelineStageCount * detail::CLCResponseSize * NumCLCResponses;
// CLC Throttle pipeline storage
static constexpr auto CLCThrottlePipelineStorage = sizeof(typename cutlass::PipelineAsync<SchedulerPipelineStageCount>::SharedStorage);
// Tmem dealloc
@@ -59,8 +61,14 @@ struct Sm100DenseGemmTmaUmmaCarveout {
static constexpr auto TmemBasePtrsStorage = SchedulerPipelineStageCount * sizeof(uint32_t);
// Tensormap Storage
static constexpr auto TensorMapStorage =
IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * NumTensorMaps /* for A and B */ :
IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * NumTensorMaps * 5 /* We have five tensormaps smem */ :
0;
// TensorMapReady pipeline storage (specific to grouped/array kernels)
static constexpr auto TensorMapReadyPipelineStorage =
IsArrayOfPointersGemm ? sizeof(typename cutlass::PipelineAsync<SchedulerPipelineStageCount>::SharedStorage) :
0;
// Smem usage that's not part of CollectiveEpilogue::SharedStorage & CollectiveMainloop::SharedStorage
static constexpr auto KernelSmemCarveout = static_cast<int>( AccumulatorPipelineStorage +
CLCPipelineStorage +
@@ -69,7 +77,8 @@ struct Sm100DenseGemmTmaUmmaCarveout {
CLCThrottlePipelineStorage +
CLCResponseStorage +
TmemBasePtrsStorage +
TensorMapStorage
TensorMapStorage +
TensorMapReadyPipelineStorage
);
};
include/cutlass/gemm/collective/builders/sm100_simt_builder.inl
+4 -1
View File
@@ -96,6 +96,7 @@ sm100_make_simt_f32_tiled_mma() {
} // namespace detail
template <
class ArchTag,
class GmemLayoutATag,
int AlignmentA,
class GmemLayoutBTag,
@@ -105,7 +106,7 @@ template <
int stages,
class BuilderScheduleTag>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassSimt,
float,
GmemLayoutATag,
@@ -119,6 +120,8 @@ struct CollectiveBuilder<
StageCount<stages>,
BuilderScheduleTag,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
) &&
(cute::is_same_v<BuilderScheduleTag, KernelMultistage> ||
cute::is_same_v<BuilderScheduleTag, KernelPtrArrayMultistage> ||
cute::is_same_v<BuilderScheduleTag, KernelScheduleAuto>) &&
include/cutlass/gemm/collective/builders/sm100_sparse_umma_builder.inl
+12 -2
View File
@@ -44,6 +44,7 @@ namespace detail {
// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count.
template <
int CapacityBytes,
class ElementAMma,
class ElementB,
class ElementEMma,
@@ -60,6 +61,7 @@ sm100_compute_stage_count_or_override_sparse(StageCount<stages> stage_count) {
// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count.
template <
int CapacityBytes,
class ElementAMma,
class ElementB,
class ElementEMma,
@@ -104,7 +106,7 @@ sm100_compute_stage_count_or_override_sparse(StageCountAutoCarveout<carveout_byt
constexpr auto EpilogueSharedStorage = carveout_bytes;
constexpr auto Stages = (cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout - EpilogueSharedStorage) /
constexpr auto Stages = (CapacityBytes - KernelSmemCarveout - EpilogueSharedStorage) /
(MainloopTensorStorage_per_Stage + MainloopPipelineStorage_per_Stage_aligned);
return Stages;
@@ -269,6 +271,7 @@ sm100_make_trivial_tiled_mma_sparse() {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ArchTag,
class ElementA,
class GmemLayoutATag,
int AlignmentA,
@@ -282,7 +285,7 @@ template <
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassSparseTensorOp,
ElementA,
GmemLayoutATag,
@@ -296,6 +299,8 @@ struct CollectiveBuilder<
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
) &&
(not cute::is_tuple_v<ElementA> && not cute::is_tuple_v<ElementB> &&
not cute::is_complex_v<ElementA> && not cute::is_complex_v<ElementB> &&
not cute::is_sparse_v<ElementA>) &&
@@ -375,7 +380,12 @@ struct CollectiveBuilder<
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
// Calculate SMEM capacity based on ArchTag
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes;
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_sparse<
ReducedSmemCapacityBytes,
ElementAMma_SmemAllocType,
ElementBMma_SmemAllocType,
ElementEMma,
include/cutlass/gemm/collective/builders/sm100_umma_builder.inl
+31 -11
View File
@@ -153,6 +153,7 @@ check_input_datatypes() {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ArchTag,
class ElementA,
class GmemLayoutATag,
int AlignmentA,
@@ -166,7 +167,7 @@ template <
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassTensorOp,
ElementA,
GmemLayoutATag,
@@ -180,6 +181,8 @@ struct CollectiveBuilder<
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
) &&
not cute::is_tuple_v<ElementA> && not cute::is_tuple_v<ElementB> &&
not cute::is_complex_v<ElementA> && not cute::is_complex_v<ElementB> &&
// Dense Gemm / PtrArrayDenseGemm
@@ -265,11 +268,17 @@ struct CollectiveBuilder<
// Calculate scheduler pipeline stages. Having one more stage than the accumulator allows more latency hiding.
using StrideA = cutlass::gemm::TagToStrideA_t<GmemLayoutATag>;
using StrideB = cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>;
using InternalStrideA = cute::remove_pointer_t<StrideA>;
using InternalStrideB = cute::remove_pointer_t<StrideB>;
static constexpr bool IsArrayOfPointersGemm = (cute::is_base_of_v<KernelScheduleSm100PtrArrayDenseGemm, BuilderScheduleTag>);
// Grouped GEMM(where Stride type is Stride*) uses specific static tile scheduler.
static constexpr bool IsGroupGemm = !cute::is_same_v<StrideA, InternalStrideA>;
static constexpr uint32_t SchedulerPipelineStageCount = cute::conditional_return<IsGroupGemm>(8, 2);
// Perform checks for both StrideA and StrideB to filter out Ragged Continguous Group Gemm
static constexpr bool IsGroupGemm = !(cute::is_same_v<StrideA, InternalStrideA>) && !(cute::is_same_v<StrideB, InternalStrideB>);
static constexpr bool IsRCGroupGemm = (cute::is_same_v<StrideA, InternalStrideA>) && !(cute::is_same_v<StrideB, InternalStrideB>);
static constexpr uint32_t SchedulerPipelineStageCount = cute::conditional_return<IsGroupGemm || IsRCGroupGemm>(8, 2);
static constexpr uint32_t KernelSmemCarveout = detail::Sm100DenseGemmTmaUmmaCarveout<
ClusterShape_MNK,
@@ -279,23 +288,34 @@ struct CollectiveBuilder<
IsArrayOfPointersGemm
>::KernelSmemCarveout;
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int Sm100ReducedSmemCapacityBytes = cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
using MainloopPipelineStorage = typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage;
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override<
Sm100ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, MainloopPipelineStorage>(StageCountType{});
ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, MainloopPipelineStorage>(StageCountType{});
static_assert(PipelineStages > 0, "Smem usage is too high. Can't create any SMEM buffers for A, and B.");
using DispatchPolicy =
cute::conditional_t<IsArrayOfPointersGemm,
cutlass::gemm::MainloopSm100ArrayTmaUmmaWarpSpecialized<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
>,
cute::conditional_t<IsRCGroupGemm,
cutlass::gemm::MainloopSm100RCGroupGemmTmaUmmaWarpSpecialized<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
>,
cutlass::gemm::MainloopSm100ArrayTmaUmmaWarpSpecialized<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
>
>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecialized<
PipelineStages,
SchedulerPipelineStageCount,
include/cutlass/gemm/collective/builders/sm103_blockscaled_umma_builder.inl
+7 -3
View File
@@ -313,6 +313,7 @@ auto sSFB = [&]() {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ArchTag,
class ElementPairA,
class GmemLayoutATag,
int AlignmentA,
@@ -326,7 +327,7 @@ template <
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm103,
ArchTag,
arch::OpClassBlockScaledTensorOp,
ElementPairA,
GmemLayoutATag,
@@ -340,6 +341,8 @@ struct CollectiveBuilder<
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm103>
) &&
// Not paired input, Not Complex input
(cute::is_tuple_v<ElementPairA> && cute::is_tuple_v<ElementPairB> &&
not cute::is_complex_v<ElementPairA> && not cute::is_complex_v<ElementPairB>) &&
@@ -495,11 +498,12 @@ struct CollectiveBuilder<
TensorMapStorage +
TmaPrefetchStorage);
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int Sm100ReducedSmemCapacityBytes = cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
using SmemTileShape = cute::Shape<Int<MMA_M>, Int<MMA_N/cute::size(AtomThrID{})>, _128>; // SmemAllocTypes are uint8_t. We always allocate 128bytes
static constexpr auto PipelineStages = cutlass::gemm::collective::detail::sm103_compute_stage_count_or_override_blockscaled<
Sm100ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, SmemLayoutAtomSFA, SmemLayoutAtomSFB>(StageCountType{});
ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, SmemLayoutAtomSFA, SmemLayoutAtomSFB>(StageCountType{});
using DispatchPolicy = typename cute::conditional_t<IsArrayOfPointersGemm,
cutlass::gemm::MainloopSm103ArrayTmaUmmaWarpSpecializedBlockScaled<
include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl
+1
View File
@@ -1140,6 +1140,7 @@ 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;
// Reserve 128B for 8 stages of tile scheduling
static constexpr size_t SchedulerPipelineStorage = cute::is_pointer_v<TagToStrideA_t<GmemLayoutATag>> ?
sizeof(cutlass::PipelineDetail::PipelineAsyncSharedStorage<8>) : 0;
include/cutlass/gemm/collective/collective_mma.hpp
+1
View File
@@ -55,6 +55,7 @@
#if !defined(__CUDACC_RTC__)
#include "cutlass/gemm/collective/sm100_mma_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_rcggemm.hpp"
#include "cutlass/gemm/collective/sm100_mma_warpspecialized_emulated.hpp"
#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp"
#include "cutlass/gemm/collective/sm100_sparse_mma_warpspecialized.hpp"
include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp
+74 -31
View File
@@ -143,6 +143,11 @@ struct CollectiveMma<
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{}))));
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{}))));
// Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously.
// This should be larger than the total number of TMA requests inflight (from update to issued to returned).
// This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies).
constexpr static uint32_t NumTmaDescriptorsPerSm = SchedulerPipelineStageCount + Stages + 2;
using ElementPairA = ElementPairA_;
using ElementPairB = ElementPairB_;
using ElementAMma = typename TiledMma::ValTypeA;
@@ -571,13 +576,18 @@ struct CollectiveMma<
};
}
struct TensorMaps : cute::aligned_struct<256, _0> {
cute::TmaDescriptor tma_desc_a;
cute::TmaDescriptor tma_desc_b;
cute::TmaDescriptor tma_desc_sfa;
cute::TmaDescriptor tma_desc_sfb;
};
template <class ProblemShape>
static size_t
get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, int sm_count) {
constexpr uint32_t NumInputTensors = 4;
constexpr size_t SizeOfCuTensorMap = sizeof(cute::TmaDescriptor);
// Allocate gmem space for input tensormaps per each SM, A tensormap copies followed by B tensormap copies
return (NumInputTensors * SizeOfCuTensorMap * sm_count);
// Allocate gmem space for input tensormaps per each SM.
return (sm_count * sizeof(TensorMaps) * NumTmaDescriptorsPerSm);
}
template <class ProblemShape>
@@ -674,7 +684,7 @@ struct CollectiveMma<
/// mcast_mask_b - tma multicast mask for B
/// mcast_mask_sfa - tma multicast mask for SFA
/// mcast_mask_sfb - tma multicast mask for SFB
template <class ProblemShape_MNKL>
template <bool IsTensorMapUpdateAsync = false, class ProblemShape_MNKL>
CUTLASS_DEVICE auto
load_init(
ProblemShape_MNKL const& problem_shape_MNKL,
@@ -682,6 +692,7 @@ struct CollectiveMma<
TensorStorage& shared_tensors,
TensorMapStorage& shared_tensormaps,
int32_t const sm_count, int32_t const sm_idx,
[[maybe_unused]] int32_t num_groups,
int32_t init_group) const {
using X = Underscore;
@@ -788,15 +799,19 @@ struct CollectiveMma<
uint16_t mcast_mask_sfa = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);
uint16_t mcast_mask_sfb = create_tma_multicast_mask<1>(cta_layout_sfb_vmnk, cta_coord_sfb_vmnk);
// Fetch a copy of tensormaps for the CTA from Params
auto input_tensormaps = tensormaps_init(params, shared_tensormaps, sm_count, sm_idx);
auto ret = cute::make_tuple(
gA_mkl, gB_nkl, // for scheduler
tAgA_mkl, tBgB_nkl, tAsA, tBsB, // for input tensor values
tAgSFA_mkl, tBgSFB_nkl, tAsSFA, tBsSFB, // for input scale factor tensor values
mcast_mask_a, mcast_mask_b, mcast_mask_sfa, mcast_mask_sfb); // multicast masks
return cute::make_tuple(
gA_mkl, gB_nkl, // for scheduler
tAgA_mkl, tBgB_nkl, tAsA, tBsB, // for input tensor values
tAgSFA_mkl, tBgSFB_nkl, tAsSFA, tBsSFB, // for input scale factor tensor values
mcast_mask_a, mcast_mask_b, mcast_mask_sfa, mcast_mask_sfb, // multicast masks
input_tensormaps); // for tma descriptor modification (per-CTA tensormap copy)
if constexpr (IsTensorMapUpdateAsync) {
return ret;
} else {
// Fetch a copy of tensormaps for the CTA from Params
auto input_tensormaps = tensormaps_init(params, shared_tensormaps, sm_count, sm_idx);
return cute::tuple_cat(ret, cute::make_tuple(input_tensormaps));
}
}
/// Set up the data needed by this collective for mma compute.
@@ -895,7 +910,8 @@ struct CollectiveMma<
cute::tuple<TensorMapA, TensorMapB, TensorMapSFA, TensorMapSFB>> const& load_inputs,
TileCoordMNKL const& cta_coord_mnkl,
KTileIterator k_tile_iter, int k_tile_count,
bool did_batch_change) {
bool did_batch_change,
[[maybe_unused]] int curr_batch) {
auto [unused_gA, unused_gB,
tAgA_mkl, tBgB_nkl, tAsA, tBsB,
@@ -1116,19 +1132,15 @@ struct CollectiveMma<
// Methods to perform different parts of TMA/Tensormap modifications
//
template <bool IsTensorMapUpdateAsync = false>
CUTLASS_DEVICE auto
tensormaps_init(
Params const& mainloop_params,
TensorMapStorage& shared_tensormaps,
int32_t const sm_count,
int32_t const sm_idx) const {
cute::TmaDescriptor* gmem_tensormap = mainloop_params.tensormaps;
cute::TmaDescriptor* tma_desc_a = &gmem_tensormap[sm_idx];
cute::TmaDescriptor* tma_desc_b = &gmem_tensormap[sm_idx + sm_count];
cute::TmaDescriptor* tma_desc_sfa = &gmem_tensormap[sm_idx + 2 * sm_count];
cute::TmaDescriptor* tma_desc_sfb = &gmem_tensormap[sm_idx + 3 * sm_count];
TensorMaps* gmem_tensormap = &(reinterpret_cast<TensorMaps*>(mainloop_params.tensormaps)[sm_idx * NumTmaDescriptorsPerSm]);
if (cute::elect_one_sync()) {
// Bringing tensormaps from params to smem for modification later
@@ -1148,9 +1160,30 @@ struct CollectiveMma<
copy(recast<uint128_t>(pSFA_tensormap), recast<uint128_t>(sSFA_tensormap));
copy(recast<uint128_t>(pSFB_tensormap), recast<uint128_t>(sSFB_tensormap));
}
__syncwarp();
return cute::make_tuple(tma_desc_a, tma_desc_b, tma_desc_sfa, tma_desc_sfb);
struct TensorMapArray {
TensorMaps *tensor_maps;
TensorMapArray() = default;
CUTLASS_DEVICE
TensorMapArray(void* tensormaps) : tensor_maps(reinterpret_cast<TensorMaps*>(tensormaps)) {}
CUTLASS_DEVICE
cute::tuple<cute::TmaDescriptor*, cute::TmaDescriptor*, cute::TmaDescriptor*, cute::TmaDescriptor*>
operator[](int32_t idx) {
idx = idx % NumTmaDescriptorsPerSm;
return cute::make_tuple(&tensor_maps[idx].tma_desc_a, &tensor_maps[idx].tma_desc_b, &tensor_maps[idx].tma_desc_sfa, &tensor_maps[idx].tma_desc_sfb);
}
};
if constexpr (IsTensorMapUpdateAsync) {
return TensorMapArray(gmem_tensormap);
} else {
return cute::make_tuple(&gmem_tensormap->tma_desc_a, &gmem_tensormap->tma_desc_b, &gmem_tensormap->tma_desc_sfa, &gmem_tensormap->tma_desc_sfb);
}
}
// Replace address for the global tensor (to be done by single thread)
@@ -1244,7 +1277,7 @@ struct CollectiveMma<
}
// The entire warp must call this function collectively (that is, the instructions are aligned)
template <class TensorMapA, class TensorMapB, class TensorMapSFA, class TensorMapSFB, class ProblemShape>
template <bool WaitForInflightTmaRequests = true, class TensorMapA, class TensorMapB, class TensorMapSFA, class TensorMapSFB, class ProblemShape>
CUTLASS_DEVICE
void
tensormaps_perform_update(
@@ -1252,10 +1285,9 @@ struct CollectiveMma<
Params const& mainloop_params,
cute::tuple<TensorMapA, TensorMapB, TensorMapSFA, TensorMapSFB> const& input_tensormaps,
ProblemShape problem_shape,
int32_t next_batch) {
int32_t next_batch
) {
if (cute::elect_one_sync()) {
// Replacing global_address for the next batch
tensormaps_replace_global_address(shared_tensormaps, mainloop_params, next_batch);
if constexpr (IsGroupedGemmKernel) {
auto problem_shape_MNKL = append<4>(problem_shape.get_problem_shape(next_batch), 1);
@@ -1263,23 +1295,34 @@ struct CollectiveMma<
tensormaps_replace_global_tensor_properties(shared_tensormaps,
mainloop_params, next_batch, problem_shape_MNKL);
}
// Replacing global_address for the next batch
tensormaps_replace_global_address(shared_tensormaps, mainloop_params, next_batch);
}
// Ensure warp is converged before issuing tensormap fence release
__syncwarp();
// Entire warp must do this (ie its aligned)
tensormaps_cp_fence_release(shared_tensormaps, input_tensormaps);
tensormaps_cp_fence_release<WaitForInflightTmaRequests>(
shared_tensormaps,
input_tensormaps
);
}
template <class TensorMapA, class TensorMapB, class TensorMapSFA, class TensorMapSFB>
template <bool WaitForInflightTmaRequests = true, class TensorMapA, class TensorMapB, class TensorMapSFA, class TensorMapSFB>
CUTLASS_DEVICE
void
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();
cute::tuple<TensorMapA, TensorMapB, TensorMapSFA, TensorMapSFB> const& input_tensormaps
) {
if constexpr (WaitForInflightTmaRequests) {
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/sm100_mma_array_warpspecialized.hpp
+63 -20
View File
@@ -123,6 +123,11 @@ struct CollectiveMma<
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{}))));
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{}))));
// Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously.
// This should be larger than the total number of TMA requests inflight (from update to issued to returned).
// This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies).
constexpr static uint32_t NumTmaDescriptorsPerSm = SchedulerPipelineStageCount + Stages + 2;
using ElementA = ElementA_;
using ElementAMma = typename TiledMma::ValTypeA;
using StrideA = StrideA_;
@@ -417,7 +422,7 @@ struct CollectiveMma<
constexpr uint32_t NumInputTensors = 2;
constexpr size_t SizeOfCuTensorMap = sizeof(cute::TmaDescriptor);
// Allocate gmem space for input tensormaps per each SM, A tensormap copies followed by B tensormap copies
return (NumInputTensors * SizeOfCuTensorMap * sm_count);
return (NumInputTensors * SizeOfCuTensorMap * sm_count * NumTmaDescriptorsPerSm);
}
template <class ProblemShape>
@@ -497,7 +502,7 @@ struct CollectiveMma<
/// tBsB - partitioned smem tensor for B
/// mcast_mask_a - tma multicast mask for A
/// mcast_mask_b - tma multicast mask for B
template <class ProblemShape_MNKL>
template <bool IsTensorMapUpdateAsync = false, class ProblemShape_MNKL>
CUTLASS_DEVICE auto
load_init(
ProblemShape_MNKL const& problem_shape_MNKL,
@@ -505,6 +510,7 @@ struct CollectiveMma<
TensorStorage& shared_tensors,
TensorMapStorage& shared_tensormaps,
int32_t const sm_count, int32_t const sm_idx,
[[maybe_unused]] int32_t num_groups,
[[maybe_unused]] int32_t init_group) const {
using X = Underscore;
@@ -550,14 +556,20 @@ struct CollectiveMma<
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);
// Fetch a copy of tensormaps for the CTA from Params
auto input_tensormaps = tensormaps_init(params, shared_tensormaps, sm_count, sm_idx);
return cute::make_tuple(
auto ret = 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
input_tensormaps); // for tma descriptor modification (per-CTA tensormap copy)
mcast_mask_a, mcast_mask_b // multicast masks
);
if constexpr (IsTensorMapUpdateAsync) {
return ret;
}
else {
// Fetch a copy of tensormaps for the CTA from Params
auto input_tensormaps = tensormaps_init<IsTensorMapUpdateAsync>(params, shared_tensormaps, sm_count, sm_idx);
return cute::tuple_cat(ret, cute::make_tuple(input_tensormaps));
}
}
/// Set up the data needed by this collective for mma compute.
@@ -612,7 +624,8 @@ struct CollectiveMma<
cute::tuple<TensorMapA, TensorMapB>> const& load_inputs,
TileCoordMNKL const& cta_coord_mnkl,
KTileIterator k_tile_iter, int k_tile_count,
bool did_batch_change) {
bool did_batch_change,
[[maybe_unused]] int curr_batch) {
auto [unused_gA, unused_gB,
tAgA_mkl, tBgB_nkl, tAsA, tBsB,
@@ -739,6 +752,7 @@ struct CollectiveMma<
// Methods to perform different parts of TMA/Tensormap modifications
//
template <bool IsTensorMapUpdateAsync = false>
CUTLASS_DEVICE auto
tensormaps_init(
Params const& mainloop_params,
@@ -747,8 +761,8 @@ struct CollectiveMma<
int32_t const sm_idx) const {
cute::TmaDescriptor* gmem_tensormap = mainloop_params.tensormaps;
cute::TmaDescriptor* tma_desc_a = &gmem_tensormap[sm_idx];
cute::TmaDescriptor* tma_desc_b = &gmem_tensormap[sm_idx + sm_count];
cute::TmaDescriptor* tma_desc_a = &gmem_tensormap[sm_idx * NumTmaDescriptorsPerSm];
cute::TmaDescriptor* tma_desc_b = &gmem_tensormap[(sm_idx + sm_count) * NumTmaDescriptorsPerSm];
if (cute::elect_one_sync()) {
// Bringing tensormaps from params to smem for modification later
@@ -762,7 +776,29 @@ struct CollectiveMma<
}
__syncwarp();
return cute::make_tuple(tma_desc_a, tma_desc_b);
struct TensorMapArray {
cute::TmaDescriptor* tma_desc_a;
cute::TmaDescriptor* tma_desc_b;
TensorMapArray() = default;
CUTLASS_DEVICE
TensorMapArray(cute::TmaDescriptor* tma_desc_a, cute::TmaDescriptor* tma_desc_b) : tma_desc_a(tma_desc_a), tma_desc_b(tma_desc_b) {}
CUTLASS_DEVICE
cute::tuple<cute::TmaDescriptor*, cute::TmaDescriptor*>
operator[](int32_t idx) {
idx = idx % NumTmaDescriptorsPerSm;
return cute::make_tuple(tma_desc_a + idx, tma_desc_b + idx);
}
};
if constexpr (IsTensorMapUpdateAsync) {
return TensorMapArray(tma_desc_a, tma_desc_b);
}
else {
return cute::make_tuple(tma_desc_a, tma_desc_b);
}
}
// Replace address for the global tensor (to be done by single thread)
@@ -826,7 +862,7 @@ struct CollectiveMma<
}
// The entire warp must call this function collectively (that is, the instructions are aligned)
template <class TensorMapA, class TensorMapB, class ProblemShape>
template <bool WaitForInflightTmaRequests = true, class TensorMapA, class TensorMapB, class ProblemShape>
CUTLASS_DEVICE
void
tensormaps_perform_update(
@@ -834,7 +870,8 @@ struct CollectiveMma<
Params const& mainloop_params,
cute::tuple<TensorMapA, TensorMapB> const& input_tensormaps,
ProblemShape problem_shape,
int32_t next_batch) {
int32_t next_batch
) {
if (cute::elect_one_sync()) {
// Replacing global_address for the next batch
tensormaps_replace_global_address(shared_tensormaps, mainloop_params, next_batch);
@@ -849,18 +886,24 @@ struct CollectiveMma<
// Ensure warp is converged before issuing tensormap fence release
__syncwarp();
// Entire warp must do this (ie its aligned)
tensormaps_cp_fence_release(shared_tensormaps, input_tensormaps);
tensormaps_cp_fence_release<WaitForInflightTmaRequests>(
shared_tensormaps,
input_tensormaps
);
}
template <class TensorMapA, class TensorMapB>
template <bool WaitForInflightTmaRequests = true, class TensorMapA, class TensorMapB>
CUTLASS_DEVICE
void
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();
cute::tuple<TensorMapA, TensorMapB> const& input_tensormaps
) {
if constexpr (WaitForInflightTmaRequests) {
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);
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_rcggemm.hpp
+899
View File
@@ -0,0 +1,899 @@
/***************************************************************************************************
* 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 "cutlass/cutlass.h"
#include "cutlass/detail/collective.hpp"
#include "cutlass/detail/cluster.hpp"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/numeric_types.h"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/gemm/gemm.h"
#include "cutlass/trace.h"
#include "cutlass/kernel_hardware_info.hpp"
#include "cutlass/cuda_host_adapter.hpp"
#include "cutlass/detail/sm100_tmem_helper.hpp"
#include "cute/algorithm/functional.hpp"
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
using namespace cute;
/////////////////////////////////////////////////////////////////////////////////////////////////
// WarpSpecialized Mainloop
// Both DMA Load and MMA methods of this class must be run by a single thread that's picked by elect_one
template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
class StrideA_,
class ElementB_,
class StrideB_,
class TiledMma_,
class GmemTiledCopyA_,
class SmemLayoutAtomA_,
class SmemCopyAtomA_,
class TransformA_,
class GmemTiledCopyB_,
class SmemLayoutAtomB_,
class SmemCopyAtomB_,
class TransformB_>
struct CollectiveMma<
MainloopSm100RCGroupGemmTmaUmmaWarpSpecialized<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
TileShape_,
ElementA_,
StrideA_,
ElementB_,
StrideB_,
TiledMma_,
GmemTiledCopyA_,
SmemLayoutAtomA_,
SmemCopyAtomA_,
TransformA_,
GmemTiledCopyB_,
SmemLayoutAtomB_,
SmemCopyAtomB_,
TransformB_>
{
//
// Type Aliases
//
using TiledMma = TiledMma_;
using AtomThrShapeMNK = Shape<decltype(shape<0>(typename TiledMma::ThrLayoutVMNK{})), _1, _1>;
using DispatchPolicy = MainloopSm100RCGroupGemmTmaUmmaWarpSpecialized<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
using TileShape = TileShape_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
CUTE_STATIC_ASSERT_V(evenly_divides(TileShape{}, tile_shape(TiledMma{})),
"Static cluster shape used: TileShape should be evenly divided by TiledMma");
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
// Define A and B block shapes for reduced size TMA_LOADs
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{}))));
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{}))));
// Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously.
// This should be larger than the total number of TMA requests inflight (from update to issued to returned).
// This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies).
constexpr static uint32_t NumTmaDescriptorsPerSm = SchedulerPipelineStageCount + Stages + 2;
using ElementA = ElementA_;
using ElementAMma = typename TiledMma::ValTypeA;
using StrideA = StrideA_;
using InternalStrideA = cute::remove_pointer_t<StrideA>;
using ElementB = ElementB_;
using ElementBMma = typename TiledMma::ValTypeB;
using StrideB = StrideB_;
using InternalStrideB = cute::remove_pointer_t<StrideB>;
static constexpr bool IsRuntimeDataTypeA = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4<ElementA>();
static constexpr bool IsRuntimeDataTypeB = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4<ElementB>();
static_assert((IsRuntimeDataTypeA && IsRuntimeDataTypeB) ||
(!IsRuntimeDataTypeA && !IsRuntimeDataTypeB),
"ElementA and ElementB should be both runtime or both static.");
static constexpr bool IsRuntimeDataType = IsRuntimeDataTypeA && IsRuntimeDataTypeB;
using ElementAccumulator = typename TiledMma::ValTypeC;
using GmemTiledCopyA = GmemTiledCopyA_;
using GmemTiledCopyB = GmemTiledCopyB_;
using SmemLayoutAtomA = SmemLayoutAtomA_;
using SmemLayoutAtomB = SmemLayoutAtomB_;
using SmemCopyAtomA = SmemCopyAtomA_;
using SmemCopyAtomB = SmemCopyAtomB_;
using TransformA = TransformA_;
using TransformB = TransformB_;
using ArchTag = typename DispatchPolicy::ArchTag;
using MainloopPipeline = cutlass::PipelineTmaUmmaAsync<
DispatchPolicy::Stages,
ClusterShape,
AtomThrShapeMNK>;
using MainloopPipelineState = typename MainloopPipeline::PipelineState;
static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtomA must be rank 2 (M,K)");
static_assert(((size<0,0>(MmaShapeA_MK{}) * size<1>(MmaShapeA_MK{})) % size<0>(SmemLayoutAtomA{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,1>(MmaShapeA_MK{}) * size<2>(MmaShapeA_MK{})) % size<1>(SmemLayoutAtomA{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(cute::is_void_v<SmemCopyAtomA>,
"SM100 UMMA cannot have a non-void copy atom for smem sourced instructions.");
static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtomB must be rank 2 (N,K)");
static_assert(((size<0,0>(MmaShapeB_NK{}) * size<1>(MmaShapeB_NK{})) % size<0>(SmemLayoutAtomB{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,1>(MmaShapeB_NK{}) * size<2>(MmaShapeB_NK{})) % size<1>(SmemLayoutAtomB{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(cute::is_void_v<SmemCopyAtomB>,
"SM100 UMMA cannot have a non-void copy atom for smem sourced instructions.");
// 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.
// (MMA_TILE_M,MMA_TILE_K),MMA_M,MMA_K,PIPE)
using SmemLayoutA = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomA{},
append(MmaShapeA_MK{}, Int<DispatchPolicy::Stages>{}),
cute::conditional_t<cutlass::gemm::detail::is_mn_major<InternalStrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
// (MMA_TILE_N,MMA_TILE_K),MMA_N,MMA_K,PIPE)
using SmemLayoutB = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomB{},
append(MmaShapeB_NK{}, Int<DispatchPolicy::Stages>{}),
cute::conditional_t<cutlass::gemm::detail::is_mn_major<InternalStrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more.");
static_assert(cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
"MMA atom must source both A and B operand from smem_desc for this mainloop.");
static_assert(
(size(AtomThrShapeMNK{}) == 1 &&
(cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>)) ||
(size(AtomThrShapeMNK{}) == 2 &&
(cute::is_same_v<GmemTiledCopyA, SM100_TMA_2SM_LOAD> || cute::is_same_v<GmemTiledCopyA, SM100_TMA_2SM_LOAD_MULTICAST>)),
"GmemTiledCopy - invalid TMA copy atom specified.");
static_assert(
(size(AtomThrShapeMNK{}) == 1 &&
(cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>)) ||
(size(AtomThrShapeMNK{}) == 2 &&
(cute::is_same_v<GmemTiledCopyB, SM100_TMA_2SM_LOAD> || cute::is_same_v<GmemTiledCopyB, SM100_TMA_2SM_LOAD_MULTICAST>)),
"GmemTiledCopy - invalid TMA copy atom specified.");
using TmaInternalElementA = cute::conditional_t<cute::is_same_v<ElementA, float>, cutlass::tfloat32_t, ElementAMma>;
using TmaInternalElementB = cute::conditional_t<cute::is_same_v<ElementB, float>, cutlass::tfloat32_t, ElementBMma>;
using SmemAllocTypeA = cute::conditional_t<cute::sizeof_bits_v<ElementAMma> < 8, uint8_t, ElementAMma>;
using SmemAllocTypeB = cute::conditional_t<cute::sizeof_bits_v<ElementBMma> < 8, uint8_t, ElementBMma>;
using BitTypeElementA = cute::uint_bit_t<cute::sizeof_bits_v<ElementA>>;
using BitTypeElementB = cute::uint_bit_t<cute::sizeof_bits_v<ElementB>>;
using ArrayElementA = cute::conditional_t<IsRuntimeDataTypeA, BitTypeElementA, ElementA>;
using ArrayElementB = cute::conditional_t<IsRuntimeDataTypeB, BitTypeElementB, ElementB>;
using RuntimeDataTypeA = cute::conditional_t<IsRuntimeDataTypeA, cute::UMMA::MXF8F6F4Format, void*>;
using RuntimeDataTypeB = cute::conditional_t<IsRuntimeDataTypeB, cute::UMMA::MXF8F6F4Format, void*>;
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128, _0> {
cute::ArrayEngine<SmemAllocTypeA, cute::cosize_v<SmemLayoutA>> smem_A;
cute::ArrayEngine<SmemAllocTypeB, cute::cosize_v<SmemLayoutB>> smem_B;
} tensors;
struct TensorMapStorage : cute::aligned_struct<128, _0> {
cute::TmaDescriptor smem_tensormap_B;
} tensormaps;
using PipelineStorage = typename MainloopPipeline::SharedStorage;
PipelineStorage pipeline;
};
// Expose shared storage for tensors/pipelines separately to allow kernel layer to reorder them.
using TensorStorage = typename SharedStorage::TensorStorage;
using TensorMapStorage = typename SharedStorage::TensorMapStorage;
using PipelineStorage = typename SharedStorage::PipelineStorage;
// Only one thread issues the TMA and updates the barriers in a 2SM MMA, adjust bytes accordingly
static constexpr uint32_t TmaTransactionBytes =
cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * cosize(take<0,3>(SmemLayoutA{})) * cute::sizeof_bits_v<ElementA>) +
cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * cosize(take<0,3>(SmemLayoutB{})) * cute::sizeof_bits_v<ElementB>);
template <class AccTensor>
struct TmemStorage {
AccTensor accumulators;
};
// Host side kernel arguments
struct Arguments {
ArrayElementA const* ptr_A{nullptr};
InternalStrideA dA{};
ArrayElementB const** ptr_B{nullptr};
StrideB dB{};
RuntimeDataTypeA runtime_data_type_a{};
RuntimeDataTypeB runtime_data_type_b{};
};
// 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<TmaInternalElementA>(
GmemTiledCopyA{},
make_tensor(recast_ptr<TmaInternalElementA>(nullptr), repeat_like(InternalStrideA{}, int32_t(0)), InternalStrideA{}),
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
ClusterLayout_VMNK{})
);
using TMA_B = decltype(make_tma_atom_B_sm100<TmaInternalElementB>(
GmemTiledCopyB{},
make_tensor(recast_ptr<TmaInternalElementB>(nullptr), repeat_like(InternalStrideB{}, int32_t(0)), InternalStrideB{}),
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;
RuntimeDataTypeA runtime_data_type_a;
RuntimeDataTypeB runtime_data_type_b;
cute::TmaDescriptor* tensormaps;
ArrayElementA const* ptr_A;
InternalStrideA dA;
ArrayElementB const** ptr_B;
StrideB dB;
};
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)
, runtime_data_type_a_(params.runtime_data_type_a)
, runtime_data_type_b_(params.runtime_data_type_b) {
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 problem_shapes,
Arguments const& args,
void* workspace,
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.
auto init_shape = repeat_like(append<4>(typename ProblemShape::UnderlyingProblemShape{}, 1), int32_t(1));
auto init_M = get<0>(init_shape);
auto init_N = get<1>(init_shape);
auto init_K_A = get<2>(init_shape);
auto init_K_B = get<2>(init_shape);
auto init_L = get<3>(init_shape);
// Tensor pointers will be fixed before the first access
auto ptr_A_first_batch = recast_ptr<TmaInternalElementA>(args.ptr_A);
TmaInternalElementB const* ptr_B_first_batch = nullptr;
auto problem_shape_MNK = problem_shapes.get_host_problem_shape(0);
init_M = get<0>(problem_shape_MNK);
init_K_A = get<2>(problem_shape_MNK);
InternalStrideA stride_a = args.dA;
InternalStrideB stride_b = InternalStrideB{};
// Batches/Groups are managed by using appropriate pointers to input matrices.
Tensor tensor_a = make_tensor(ptr_A_first_batch, make_layout(make_shape(init_M, init_K_A, problem_shapes.groups()), stride_a));
Tensor tensor_b = make_tensor(ptr_B_first_batch, make_layout(make_shape(init_N, init_K_B, init_L), stride_b));
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);
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<TmaInternalElementA>(
GmemTiledCopyA{},
tensor_a,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_B tma_load_b = make_tma_atom_B_sm100<TmaInternalElementB>(
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<TmaInternalElementA>(
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<TmaInternalElementB>(
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,
args.runtime_data_type_a,
args.runtime_data_type_b,
reinterpret_cast<cute::TmaDescriptor*>(workspace),
args.ptr_A,
args.dA,
reinterpret_cast<ArrayElementB const**>(args.ptr_B),
args.dB
};
}
template <class ProblemShape>
static size_t
get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, int sm_count) {
constexpr uint32_t NumInputTensors = 1;
constexpr size_t SizeOfCuTensorMap = sizeof(cute::TmaDescriptor);
// Allocate gmem space for input tensormaps per each SM, B tensormap copies
return (NumInputTensors * SizeOfCuTensorMap * sm_count * NumTmaDescriptorsPerSm);
}
template <class ProblemShape>
static cutlass::Status
initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, CudaHostAdapter* cuda_adapter = nullptr) {
return cutlass::Status::kSuccess;
}
template <class ProblemShape>
static bool
can_implement(
ProblemShape problem_shapes,
[[maybe_unused]] Arguments const& args) {
static constexpr bool IsF8F6F4 = detail::is_sm100_mma_f8f6f4<TiledMma, ElementA, ElementB>();
constexpr int tma_alignment_bits_A = cutlass::detail::get_input_alignment_bits<ElementA, IsF8F6F4>();
constexpr int tma_alignment_bits_B = cutlass::detail::get_input_alignment_bits<ElementB, IsF8F6F4>();
constexpr int min_tma_aligned_elements_A = tma_alignment_bits_A / cute::sizeof_bits<ElementA>::value;
constexpr int min_tma_aligned_elements_B = tma_alignment_bits_B / cute::sizeof_bits<ElementB>::value;
bool implementable = true;
if (problem_shapes.is_host_problem_shape_available()) {
// Check alignment for all problem sizes
for (int i = 0; i < problem_shapes.groups(); i++) {
auto problem_shape_MNKL = append<4>(problem_shapes.get_host_problem_shape(i), 1);
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{});
}
}
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
}
return implementable;
}
/// Construct A Single Stage's Accumulator Shape
CUTLASS_DEVICE static
auto
partition_accumulator_shape() {
return partition_shape_C(TiledMma{}, take<0,2>(TileShape{})); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N)
}
template <class TmemStorage>
CUTLASS_DEVICE static
auto
slice_accumulator(TmemStorage tmem_storage, int stage) {
return tmem_storage.accumulators(_,_,_,stage);
}
template <class EpilogueTile, bool IsOverlappingAccum = false>
CUTLASS_DEVICE static
auto
init_tmem_tensors(EpilogueTile epi_tile) {
TiledMma tiled_mma;
auto acc_shape = partition_accumulator_shape();
// ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,ACC_PIPE) where ACC_PIPE=2 so we can double buffer our accumulators for mainloop and epilogue.
Tensor accumulators = cutlass::detail::make_sm100_accumulator<AccumulatorPipelineStageCount, IsOverlappingAccum>(
tiled_mma, acc_shape, EpilogueTile{});
TmemStorage<decltype(accumulators)> tmem_storage;
tmem_storage.accumulators = accumulators;
return tmem_storage;
}
template <class TmemStorage>
CUTLASS_DEVICE static
void
set_tmem_offsets(TmemStorage& tmem_storage, uint32_t tmem_base_addr) {
tmem_storage.accumulators.data() = tmem_base_addr;
}
/// Set up the data needed by this collective for load.
/// Return tuple element contain
/// gA_mkl - The tiled tma tensor for input A
/// gB_nkl - The tiled tma tensor for input B
/// tAsA - partitioned smem tensor for A
/// tBsB - partitioned smem tensor for B
/// mcast_mask_a - tma multicast mask for A
/// mcast_mask_b - tma multicast mask for B
template <bool IsTensorMapUpdateAsync = false, class ProblemShape_MNKL>
CUTLASS_DEVICE auto
load_init(
ProblemShape_MNKL const& problem_shape_MNKL,
Params const& params,
TensorStorage& shared_tensors,
TensorMapStorage& shared_tensormaps,
int32_t const sm_count, int32_t const sm_idx,
int32_t num_groups,
[[maybe_unused]] int32_t init_group) const {
using X = Underscore;
// Separate out problem shape for convenience
auto [M,N,K,L] = problem_shape_MNKL;
// Problem Shape and therefore strides that we construct are [M,N,K,L], but since here for the TMA loads
// we are managing TMA descriptors to change batches, we need to neglect the L mode
const int32_t mock_L = 1;
// Represent the full tensors -- get these from TMA
Tensor mA_mkl = observed_tma_load_a_->get_tma_tensor(make_shape(M,K,num_groups));
Tensor mB_nkl = observed_tma_load_b_->get_tma_tensor(make_shape(N,K,mock_L));
// Tile the tensors and defer the slice
Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M, BLK_K, m, k, l)
Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l)
// Partition for this CTA
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_tensors.smem_A.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.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);
auto ret = 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
);
if constexpr (IsTensorMapUpdateAsync) {
return ret;
}
else {
// Fetch a copy of tensormaps for the CTA from Params
auto input_tensormaps = tensormaps_init<IsTensorMapUpdateAsync>(params, shared_tensormaps, sm_count, sm_idx);
return cute::tuple_cat(ret, cute::make_tuple(input_tensormaps));
}
}
/// Set up the data needed by this collective for mma compute.
template <class TmemStorage>
CUTLASS_DEVICE auto
mma_init(
[[maybe_unused]] TmemStorage tmem_storage,
TensorStorage& shared_tensors) const {
// Allocate "fragments/descriptors" for A and B matrices
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
// Allocate "fragments/descriptors" for A and B matrices
Tensor tCrA = TiledMma::make_fragment_A(sA); // (MMA,MMA_M,MMA_K,PIPE)
Tensor tCrB = TiledMma::make_fragment_B(sB); // (MMA,MMA_N,MMA_K,PIPE)
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<3>(sA)); // PIPE
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<3>(sB)); // PIPE
TiledMma tiled_mma;
if constexpr (IsRuntimeDataType) {
// Update instruction descriptor according to runtime argument.
// Applying bitmask (0b111) to help compiler deduce that the conversion and assignment are safe.
tiled_mma.idesc_.a_format_ = uint8_t(runtime_data_type_a_) & 0b111;
tiled_mma.idesc_.b_format_ = uint8_t(runtime_data_type_b_) & 0b111;
}
return cute::make_tuple(tiled_mma, tCrA, tCrB);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Producer Perspective
template <
class GTensorA, class GTensorB,
class GTensorPartitionedA, class GTensorPartitionedB,
class STensorA, class STensorB,
class TensorMapB,
class TileCoordMNKL,
class KTileIterator
>
CUTLASS_DEVICE auto
load(
Params const& params,
MainloopPipeline mainloop_pipeline,
MainloopPipelineState mainloop_pipe_producer_state,
cute::tuple<GTensorA, GTensorB,
GTensorPartitionedA, GTensorPartitionedB,
STensorA, STensorB,
uint16_t, uint16_t,
cute::tuple<TensorMapB>> const& load_inputs,
TileCoordMNKL const& cta_coord_mnkl,
KTileIterator k_tile_iter, int k_tile_count,
bool did_batch_change,
int curr_batch) {
auto [unused_gA, unused_gB,
tAgA_mkl, tBgB_nkl, tAsA, tBsB,
mcast_mask_a, mcast_mask_b,
input_tensormaps] = load_inputs;
// Check to see if tensormaps have been replaced in gmem
if (did_batch_change) {
tensormaps_fence_acquire(input_tensormaps);
}
// slice out the work coord from partitioned tensors
Tensor tAgA = tAgA_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, curr_batch);
Tensor tBgB = tBgB_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
auto barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state);
// Issue the Mainloop loads
CUTLASS_PRAGMA_NO_UNROLL
while (k_tile_count > 0) {
// LOCK mainloop_pipe_producer_state for _writing_
mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state, barrier_token);
using BarrierType = typename MainloopPipeline::ProducerBarrierType;
BarrierType* tma_barrier = mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state);
int write_stage = mainloop_pipe_producer_state.index();
++mainloop_pipe_producer_state;
barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state);
if (cute::elect_one_sync()) {
copy(observed_tma_load_a_->with(*tma_barrier, mcast_mask_a), tAgA(_,*k_tile_iter), tAsA(_,write_stage));
copy(observed_tma_load_b_->with(get<0>(input_tensormaps), *tma_barrier, mcast_mask_b), tBgB(_,*k_tile_iter), tBsB(_,write_stage));
}
--k_tile_count;
++k_tile_iter;
}
return cute::make_tuple(mainloop_pipe_producer_state, k_tile_iter);
}
/// Perform a Producer Epilogue to prevent early exit of ctas in a Cluster
CUTLASS_DEVICE void
load_tail(MainloopPipeline mainloop_pipeline, MainloopPipelineState mainloop_pipe_producer_state) {
// Issue the epilogue waits
// This helps avoid early exit of ctas in Cluster
// Waits for all stages to either be released (all
// Consumer UNLOCKs), or if the stage was never used
// then would just be acquired since the phase was
// still inverted from make_producer_start_state
mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class AccumulatorPipeline,
class FrgEngine, class FrgLayout,
class FragmentA, class FragmentB,
class CtaTileCoord
>
CUTLASS_DEVICE auto
mma(cute::tuple<MainloopPipeline,
AccumulatorPipeline> pipelines,
cute::tuple<MainloopPipelineState,
typename AccumulatorPipeline::PipelineState> pipeline_states,
cute::Tensor<FrgEngine, FrgLayout>& accumulators,
cute::tuple<TiledMma, FragmentA, FragmentB> const& mma_inputs,
CtaTileCoord cta_tile_coord,
int k_tile_count
) {
static_assert(is_tmem<FrgEngine>::value, "Accumulator must be tmem resident.");
static_assert(rank(FrgLayout{}) == 3, "Accumulator must be MMA-partitioned: (MMA, MMA_M, MMA_N)");
auto [tiled_mma, tCrA, tCrB] = mma_inputs;
auto [mainloop_pipeline, accumulator_pipeline] = pipelines;
auto [mainloop_pipe_consumer_state, accumulator_pipe_producer_state] = pipeline_states;
uint32_t skip_wait = k_tile_count <= 0;
auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);
//
// PIPELINED MAIN LOOP
//
tiled_mma.accumulate_ = UMMA::ScaleOut::Zero;
// Wait for tmem accumulator buffer to become empty with a flipped phase
accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);
CUTLASS_PRAGMA_NO_UNROLL
while (k_tile_count > 0) {
// WAIT on mainloop_pipe_consumer_state until its data are available
// (phase bit flips from mainloop_pipe_consumer_state.phase() value)
mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);
// Compute on k_tile
int read_stage = mainloop_pipe_consumer_state.index();
// Save current mainlop pipeline read state
auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;
// Advance mainloop_pipe
++mainloop_pipe_consumer_state;
--k_tile_count;
skip_wait = k_tile_count <= 0;
// Peek at next iteration
barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);
// Unroll the K mode manually so we can set scale C 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),
accumulators);
tiled_mma.accumulate_ = UMMA::ScaleOut::One;
}
mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state);
}
return mainloop_pipe_consumer_state;
}
//
// Methods to perform different parts of TMA/Tensormap modifications
//
template <bool IsTensorMapUpdateAsync = false>
CUTLASS_DEVICE auto
tensormaps_init(
Params const& mainloop_params,
TensorMapStorage& shared_tensormaps,
int32_t const sm_count,
int32_t const sm_idx) const {
cute::TmaDescriptor* gmem_tensormap = mainloop_params.tensormaps;
cute::TmaDescriptor* tma_desc_b = &gmem_tensormap[sm_idx * NumTmaDescriptorsPerSm];
if (cute::elect_one_sync()) {
// Bringing tensormaps from params to smem for modification later
Tensor pB_tensormap = make_tensor(observed_tma_load_b_->get_tma_descriptor(), Int<1>{}, Int<1>{});
Tensor sB_tensormap = make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_B), Int<1>{}, Int<1>{});
copy(recast<uint128_t>(pB_tensormap), recast<uint128_t>(sB_tensormap));
}
__syncwarp();
struct TensorMapArray {
cute::TmaDescriptor* tma_desc_b;
TensorMapArray() = default;
CUTLASS_DEVICE
TensorMapArray(cute::TmaDescriptor* tma_desc_b) : tma_desc_b(tma_desc_b) {}
CUTLASS_DEVICE
cute::tuple<cute::TmaDescriptor*>
operator[](int32_t idx) {
idx = idx % NumTmaDescriptorsPerSm;
return cute::make_tuple(tma_desc_b + idx);
}
};
if constexpr (IsTensorMapUpdateAsync) {
return TensorMapArray(tma_desc_b);
}
else {
return cute::make_tuple(tma_desc_b);
}
}
// Replace address for the global tensor (to be done by single thread)
CUTLASS_DEVICE
void
tensormaps_replace_global_address(
TensorMapStorage& shared_tensormaps,
Params const& mainloop_params,
int32_t next_batch) {
// Replacing global_address for the next batch
cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormaps.smem_tensormap_B,
mainloop_params.ptr_B[next_batch]);
}
// Replace dim and strides for the global tensor - used only for Grouped GEMM (to be done by single thread)
template <class ProblemShape_MNKL>
CUTLASS_DEVICE
void
tensormaps_replace_global_tensor_properties(
TensorMapStorage& shared_tensormaps,
Params const& mainloop_params,
int32_t next_group,
ProblemShape_MNKL problem_shape_mnkl) {
const uint32_t N = get<1>(problem_shape_mnkl);
const uint32_t K = get<2>(problem_shape_mnkl);
// Replace all dims for consistency
constexpr int MaxTensorRank = 5;
cute::array<uint32_t, MaxTensorRank> prob_shape_B = {1,1,1,1,1};
cute::array<uint64_t, MaxTensorRank> prob_stride_B = {0,0,0,0,0};
TmaInternalElementB const* ptr_B = nullptr;
Tensor tensor_b = make_tensor(ptr_B, make_shape(N,K,Int<1>{}), mainloop_params.dB[next_group]);
cute::detail::fill_tma_gmem_shape_stride(*observed_tma_load_b_, tensor_b,
prob_shape_B, prob_stride_B);
// Convert strides to byte strides
for (uint64_t& stride : prob_stride_B) {
stride = (stride * sizeof_bits_v<TmaInternalElementB>) / 8;
}
cute::tma_descriptor_replace_dims_strides_in_shared_mem(shared_tensormaps.smem_tensormap_B,
prob_shape_B,
prob_stride_B);
}
// The entire warp must call this function collectively (that is, the instructions are aligned)
template <bool WaitForInflightTmaRequests = true, class TensorMapB, class ProblemShape>
CUTLASS_DEVICE
void
tensormaps_perform_update(
TensorMapStorage& shared_tensormaps,
Params const& mainloop_params,
cute::tuple<TensorMapB> const& input_tensormaps,
ProblemShape problem_shape,
int32_t next_batch
) {
if (cute::elect_one_sync()) {
// Replacing global_address for the next batch
tensormaps_replace_global_address(shared_tensormaps, mainloop_params, next_batch);
auto problem_shape_MNKL = append<4>(problem_shape.get_problem_shape(next_batch), 1);
// Replacing global dims and strides for the next batch
tensormaps_replace_global_tensor_properties(shared_tensormaps,
mainloop_params, next_batch, problem_shape_MNKL);
}
// Ensure warp is converged before issuing tensormap fence release
__syncwarp();
// Entire warp must do this (ie its aligned)
tensormaps_cp_fence_release<WaitForInflightTmaRequests>(
shared_tensormaps,
input_tensormaps
);
}
template <bool WaitForInflightTmaRequests = true, class TensorMapB>
CUTLASS_DEVICE
void
tensormaps_cp_fence_release (
TensorMapStorage& shared_tensormaps,
cute::tuple<TensorMapB> const& input_tensormaps
) {
if constexpr (WaitForInflightTmaRequests) {
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_B);
}
// The entire warp must call this function collectively (that is, the instructions are aligned)
template <class TensorMapB>
CUTLASS_DEVICE
void
tensormaps_fence_acquire(cute::tuple<TensorMapB> const& input_tensormaps) {
cute::tma_descriptor_fence_acquire(get<0>(input_tensormaps));
}
protected:
typename Params::TMA_A const* observed_tma_load_a_{nullptr};
typename Params::TMA_B const* observed_tma_load_b_{nullptr};
RuntimeDataTypeA runtime_data_type_a_{};
RuntimeDataTypeB runtime_data_type_b_{};
ClusterShape cluster_shape_;
uint32_t block_rank_in_cluster_;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/sm100_mma_warpspecialized_mixed_input.hpp
+29 -31
View File
@@ -262,7 +262,6 @@ public:
AtomThrShapeMNK>;
using Mma2AccumPipelineState = typename Mma2AccumPipeline::PipelineState;
static constexpr int ScaleGranularityMN = size<0,0>(LayoutScale{});
static constexpr int ScaleGranularityK = size<1,0>(LayoutScale{});
using ScaleConfig = cutlass::detail::Sm100MixedInputBlockwiseScaleConfig<
@@ -273,11 +272,10 @@ public:
decltype(make_shape(size<0>(TileShape{}), size<2>(TileShape{}))),
decltype(make_shape(size<1>(TileShape{}), size<2>(TileShape{})))>;
static constexpr int ScaleTileShape_MN = get<0>(ScaleTileShape{});
using SmemLayoutAtomScaleFull = decltype(ScaleConfig::smem_atom_layout_scale(ScaleTileShape{}));
static constexpr int ScaleK = get<1>(ScaleTileShape{}) / ScaleGranularityK;
using SmemLayoutAtomScale = decltype(ScaleConfig::smem_atom_layout_scale(ScaleTileShape{}));
// Getting the SmemSizeMN and SmemSizeK from the mixed_dtype blockwise utils.
using SmemLayoutAtomScale = decltype(slice(make_coord(make_coord(_,0),make_coord(_,0)), SmemLayoutAtomScaleFull{}));
static_assert(cute::rank(InternalSmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
static_assert((size<0>(TileShape{}) % size<0>(InternalSmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
@@ -324,10 +322,10 @@ public:
append(CtaShapeB_NK{}, Int<DispatchPolicy::Load2TransformPipelineStageCount>{}),
(cute::conditional_t<cutlass::gemm::detail::is_mn_major<StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})));
using SmemLayoutScale = decltype(make_layout(
append(shape(SmemLayoutAtomScale{}), Int<DispatchPolicy::Load2TransformPipelineStageCount>{}),
append(stride(SmemLayoutAtomScale{}), size(filter_zeros(SmemLayoutAtomScale{})))
));
using SmemLayoutScale = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomScale{},
append(CtaShapeA_MK{}, Int<DispatchPolicy::Load2TransformPipelineStageCount>{}),
(cute::conditional_t<cutlass::gemm::detail::is_mn_major<StrideA>(), 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.");
@@ -437,12 +435,13 @@ public:
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_Scale = decltype(make_tma_atom(
using TMA_Scale = decltype(make_tma_atom_A_sm100(
GmemTiledCopyScale{},
make_tensor(static_cast<NonVoidElementScale const*>(nullptr), LayoutScale{}),
SmemLayoutScale{}(_,_,cute::Int<0>{}),
ScaleTileShape{},
size<2>(ClusterLayout_VMNK{}))
SmemLayoutScale{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
ClusterLayout_VMNK{})
);
TMA_Scale tma_load_scale;
@@ -576,13 +575,13 @@ public:
ElementScale const* ptr_S = args.ptr_S;
Tensor tensor_scale = make_tensor(detail::get_logical_ptr(ptr_S), args.layout_S);
typename Params::TMA_Scale tma_load_scale = make_tma_atom(
GmemTiledCopyScale{},
tensor_scale,
SmemLayoutScale{}(_,_,cute::Int<0>{}),
ScaleTileShape{},
size<2>(cluster_layout_vmnk)
);
typename Params::TMA_Scale tma_load_scale = make_tma_atom_A_sm100<ElementScale>(
GmemTiledCopyScale{},
tensor_scale,
SmemLayoutScale{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
if constexpr(KernelConversionMode == ConversionMode::ConvertAndScale) {
typename Params::TMAScaleParams scale_params{tma_load_scale, {}};
@@ -598,12 +597,13 @@ public:
}
else if constexpr(KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
Tensor tensor_zero = make_tensor(detail::get_logical_ptr(args.ptr_Z), args.layout_S);
typename Params::TMA_Scale tma_load_zero = make_tma_atom(
GmemTiledCopyScale{},
tensor_zero,
SmemLayoutScale{}(_,_,cute::Int<0>{}),
ScaleTileShape{},
size<2>(cluster_layout_vmnk));
typename Params::TMA_Scale tma_load_zero = make_tma_atom_A_sm100<ElementScale>(
GmemTiledCopyScale{},
tensor_zero,
SmemLayoutScale{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMAScaleParams scale_params{tma_load_scale, tma_load_zero};
return {
@@ -932,12 +932,11 @@ public:
Tensor sS = make_tensor(make_smem_ptr(shared_storage.input.smem_scale.begin()), SmemLayoutScale{});
Tensor tCgS_mkl = cta_mma.partition_A(gS_mkl); // (MMA, MMA_M, MMA_K, m, k, l)
Tensor tCsS = cta_mma.partition_A(sS);
// Project the cta_layout for tma_scale along the n-modes
auto [tSgS_mkl, tSsS] = tma_partition(params.tma_load_scale,
get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),
group_modes<0,3>(tCsS), group_modes<0,3>(tCgS_mkl));
group_modes<0,3>(sS), group_modes<0,3>(tCgS_mkl));
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
return cute::make_tuple(
@@ -953,11 +952,10 @@ public:
Tensor tCgZ_mkl = cta_mma.partition_A(gZ_mkl); // (MMA, MMA_M, MMA_K, m, k, l)
Tensor tCsZ = cta_mma.partition_A(sZ);
// Project the cta_layout for tma_scale along the n-modes
auto [tZgZ_mkl, tZsZ] = tma_partition(params.tma_load_zero,
get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),
group_modes<0,3>(tCsZ), group_modes<0,3>(tCgZ_mkl));
group_modes<0,3>(sZ), group_modes<0,3>(tCgZ_mkl));
return cute::make_tuple(
gA_mkl, gB_nkl, // for scheduler
tAgA_mkl, tBgB_nkl, tAsA, tBsB, // for input tensor values
@@ -1134,7 +1132,7 @@ public:
setup_copy_ops(sA, InputCopyAtomA{}, sACompute, [&](auto &arg) {return TiledMma::make_fragment_A(arg);}, ComputeCopyAtomA{});
// Partition of thread -> shared and thread -> RF
auto fragment_compute = TiledMma::make_fragment_A(sACompute);
auto fragment_compute = TiledMma::make_fragment_A(sS);
fragment_compute.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators);
auto r2t_tiled_copy = make_tmem_copy(ComputeCopyAtomA{}, fragment_compute(_,_,_,0));
auto src_copy_scale = make_tiled_copy_S(Copy_Atom<DefaultCopy, ElementScale>{}, r2t_tiled_copy);
include/cutlass/gemm/collective/sm120_blockscaled_sparse_mma_tma.hpp
+2 -1
View File
@@ -101,7 +101,8 @@ struct CollectiveMma<
using StridePairB = StridePairB_;
using SmemCopyAtomsA = SmemCopyAtomsA_;
using SmemCopyAtomsB = SmemCopyAtomsB_;
using RuntimeDataTypeA = void*;
using RuntimeDataTypeB = void*;
using TiledMma = TiledMma_;
using AtomThrShapeMNK = Shape<decltype(shape<0>(typename TiledMma::ThrLayoutVMNK{})), _1, _1>;
using DispatchPolicy = MainloopSm120TmaWarpSpecializedSparseBlockScaled<StagesA, StagesB, StagesE, SchedulerPipelineStageCount, ClusterShape>;
include/cutlass/gemm/dispatch_policy.hpp
+16 -2
View File
@@ -539,7 +539,7 @@ struct KernelTmaWarpSpecializedInputTransformSm100 final {
static constexpr int AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;
};
// InputTransform GEMM
// Mixed Input Transform GEMM
template<
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_
@@ -1177,6 +1177,21 @@ struct MainloopSm100ArrayTmaUmmaWarpSpecialized {
using Schedule = KernelPtrArrayTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
>
struct MainloopSm100RCGroupGemmTmaUmmaWarpSpecialized {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
constexpr static bool IsOverlappingAccum = false;
using Schedule = KernelPtrArrayTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
template<
int Stages_,
@@ -1241,7 +1256,6 @@ struct MainloopSm100ArrayTmaUmmaWarpSpecializedFastF32 {
constexpr static int Stages = Load2TransformPipelineStageCount;
};
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
template<
int LoadABPipelineStageCount_,
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp
+497 -89
View File
@@ -132,7 +132,10 @@ public:
using AtomThrShapeMNK = typename CollectiveMainloop::AtomThrShapeMNK;
using CtaShape_MNK = typename CollectiveMainloop::CtaShape_MNK;
static constexpr bool IsGroupedGemmKernel = !cute::is_same_v<InternalStrideA, StrideA>;
//For the case of RCGroupedGemm we are still GroupedGemm but our StrideA will not match with InternalStrideA
// Hence it's better to take this decision based upon StrideB
static constexpr bool IsGroupedGemmKernel = !(cute::is_same_v<StrideB, InternalStrideB>);
using TileSchedulerTag = cute::conditional_t<IsGroupedGemmKernel, GroupScheduler, TileSchedulerTag_>;
using TileScheduler = typename detail::TileSchedulerSelector<
@@ -140,25 +143,40 @@ public:
using TileSchedulerArguments = typename TileScheduler::Arguments;
using TileSchedulerParams = typename TileScheduler::Params;
static constexpr bool IsSchedDynamicPersistent = TileScheduler::IsDynamicPersistent;
static constexpr bool IsTensorMapUpdateAsync = not IsSchedDynamicPersistent;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
static constexpr uint32_t MinTensorMapWorkspaceAlignment = 64;
// Warp specialization thread count per threadblock
static constexpr uint32_t NumSchedThreads = NumThreadsPerWarp; // 1 warp
static constexpr uint32_t NumMMAThreads = NumThreadsPerWarp; // 1 warp
static constexpr uint32_t NumMainloopLoadThreads = NumThreadsPerWarp; // 1 warp
static constexpr uint32_t NumEpilogueLoadThreads = NumThreadsPerWarp; // 1 warp
static constexpr uint32_t NumEpilogueThreads = CollectiveEpilogue::ThreadCount;
static constexpr uint32_t NumEpilogueWarps = NumEpilogueThreads / NumThreadsPerWarp;
static constexpr uint32_t NumSchedThreads = NumThreadsPerWarp; // 1 warp
static constexpr uint32_t NumTensorMapUpdaterThreads = IsTensorMapUpdateAsync ? NumThreadsPerWarp * 4 : 0; // Four warps to update tensor maps and plumb updated tileId.
static constexpr uint32_t NumMMAThreads = NumThreadsPerWarp; // 1 warp
static constexpr uint32_t NumMainloopLoadThreads = NumThreadsPerWarp; // 1 warp
static constexpr uint32_t NumEpilogueLoadThreads = NumThreadsPerWarp; // 1 warp
static constexpr uint32_t NumEpilogueThreads = CollectiveEpilogue::ThreadCount;
static constexpr uint32_t NumEpilogueWarps = NumEpilogueThreads / NumThreadsPerWarp;
static constexpr uint32_t MaxThreadsPerBlock = NumSchedThreads +
static_assert(
SchedulerPipelineStageCount % (IsTensorMapUpdateAsync ? NumTensorMapUpdaterThreads / NumThreadsPerWarp : 1) == 0,
"SchedulerPipelineStageCount for async tensor map update kernels must be divisible by the number of asynchronous tensor map updater warps."
);
static_assert(
(!IsTensorMapUpdateAsync)
|| CollectiveEpilogue::NumMaxSchedulerPipelineStageCount >= SchedulerPipelineStageCount,
"The epilog collective expected a less scheduler stage count. Consider relaxing its NumMaxSchedulerPipelineStageCount parameter."
);
static constexpr uint32_t MaxThreadsPerBlock = NumSchedThreads + NumTensorMapUpdaterThreads +
NumMainloopLoadThreads + NumMMAThreads +
NumEpilogueLoadThreads + NumEpilogueThreads;
static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
static constexpr uint32_t NumFixupBarriers = 1;
static constexpr uint32_t CLCResponseSize = sizeof(typename TileScheduler::CLCResponse);
static constexpr bool IsSchedDynamicPersistent = TileScheduler::IsDynamicPersistent;
static constexpr uint32_t GenericRegisterRequirement = 136;
static constexpr uint32_t AccumRegisterRequirement = 232;
// Pipeline and pipeline state types
using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
@@ -179,17 +197,42 @@ public:
cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape>,
cutlass::PipelineAsync<SchedulerPipelineStageCount>>;
using CLCPipelineState = typename CLCPipeline::PipelineState;
using TensorMapReadyPipeline = cute::conditional_t<IsTensorMapUpdateAsync,
cutlass::PipelineAsync<SchedulerPipelineStageCount>,
CLCPipeline
>;
using TensorMapReadyPipelineState = typename TensorMapReadyPipeline::PipelineState;
using CLCThrottlePipeline = cute::conditional_t<IsSchedDynamicPersistent,
cutlass::PipelineAsync<SchedulerPipelineStageCount>,
cutlass::PipelineEmpty>;
using CLCThrottlePipelineState = typename CLCThrottlePipeline::PipelineState;
template <class BaseResponse>
struct WithTensorMapUpdateInfo : public BaseResponse {
uint16_t batch_changed = 0;
uint16_t TMA_stage = 0;
WithTensorMapUpdateInfo() = default;
CUTLASS_DEVICE WithTensorMapUpdateInfo(BaseResponse const& response) : BaseResponse(response) {}
};
using CLCResponseWithAdditionalInformation = cute::conditional_t<
IsTensorMapUpdateAsync,
WithTensorMapUpdateInfo<typename TileScheduler::CLCResponse>,
typename TileScheduler::CLCResponse
>;
using TmemAllocator = cute::conditional_t<cute::size(cute::shape<0>(typename TiledMma::ThrLayoutVMNK{})) == 1,
cute::TMEM::Allocator1Sm, cute::TMEM::Allocator2Sm>;
// Kernel level shared memory storage
struct SharedStorage {
struct PipelineStorage : cute::aligned_struct<16, _1> {
// The PipelineStorageImplWithoutAsyncUpdate and PipelineStorageImplWithAsyncUpdate only differ in the
// presence of the TensorMapReadyPipelineStorage.
// We could use some other technique to avoid duplication for the common members, but any technique
// we tried would break the MSVC build.
// As a workaround, we just copied the code.
struct PipelineStorageImplWithoutAsyncUpdate : cute::aligned_struct<16, _1> {
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
using LoadOrderBarrierStorage = typename LoadOrderBarrier::SharedStorage;
@@ -204,9 +247,34 @@ public:
alignas(16) AccumulatorPipelineStorage accumulator;
alignas(16) CLCThrottlePipelineStorage clc_throttle;
alignas(16) arch::ClusterBarrier tmem_dealloc;
} pipelines;
};
alignas(16) typename TileScheduler::CLCResponse clc_response[SchedulerPipelineStageCount];
struct PipelineStorageImplWithAsyncUpdate : cute::aligned_struct<16, _1> {
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
using LoadOrderBarrierStorage = typename LoadOrderBarrier::SharedStorage;
using CLCPipelineStorage = typename CLCPipeline::SharedStorage;
using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;
using CLCThrottlePipelineStorage = typename CLCThrottlePipeline::SharedStorage;
alignas(16) MainloopPipelineStorage mainloop;
alignas(16) EpiLoadPipelineStorage epi_load;
alignas(16) LoadOrderBarrierStorage load_order;
alignas(16) CLCPipelineStorage clc;
alignas(16) AccumulatorPipelineStorage accumulator;
alignas(16) CLCThrottlePipelineStorage clc_throttle;
alignas(16) arch::ClusterBarrier tmem_dealloc;
// Below is the only difference between PipelineStorageImpl<false> and PipelineStorageImpl<true>
using TensorMapReadyPipelineStorage = typename TensorMapReadyPipeline::SharedStorage;
alignas(16) TensorMapReadyPipelineStorage tensor_map_ready;
};
using PipelineStorage = cute::conditional_t<IsTensorMapUpdateAsync, PipelineStorageImplWithAsyncUpdate, PipelineStorageImplWithoutAsyncUpdate>;
PipelineStorage pipelines;
alignas(16) CLCResponseWithAdditionalInformation clc_response[IsTensorMapUpdateAsync ? 2 : 1][SchedulerPipelineStageCount];
uint32_t tmem_base_ptr;
struct TensorMapStorage : cute::aligned_struct<128, _1> {
@@ -214,7 +282,7 @@ public:
using MainloopTensorMapStorage = typename CollectiveMainloop::TensorMapStorage;
alignas(128) EpilogueTensorMapStorage epilogue;
alignas(128) MainloopTensorMapStorage mainloop;
} tensormaps;
} tensormaps[(NumTensorMapUpdaterThreads/NumThreadsPerWarp)+1];
struct TensorStorage : cute::aligned_struct<128, _1> {
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
@@ -226,7 +294,6 @@ public:
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Host facing host arguments
struct Arguments {
@@ -253,7 +320,9 @@ public:
Sched = 1,
MainloopLoad = 2,
EpilogueLoad = 3,
Epilogue = 4
Epilogue = 4,
// TensorMapUpdater starts at 256 thread alignment
TensorMapUpdater = 8
};
struct IsParticipant {
@@ -262,6 +331,7 @@ public:
uint32_t main_load = false;
uint32_t epi_load = false;
uint32_t epilogue = false;
uint32_t tensor_map_updater = false;
};
//
@@ -480,18 +550,20 @@ public:
using namespace cute;
using X = Underscore;
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
auto problem_shape = params.problem_shape;
// Account for more than one epilogue warp
int warp_idx = canonical_warp_idx_sync();
WarpCategory warp_category = warp_idx < static_cast<int>(WarpCategory::Epilogue) ? WarpCategory(warp_idx)
: WarpCategory::Epilogue;
WarpCategory warp_category = warp_idx < static_cast<int>(WarpCategory::Epilogue) ? WarpCategory(warp_idx)
: warp_idx < static_cast<int>(WarpCategory::TensorMapUpdater) ? WarpCategory::Epilogue
: WarpCategory::TensorMapUpdater;
uint32_t lane_predicate = cute::elect_one_sync();
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{});
int cluster_size = size(cluster_shape);
uint32_t cta_rank_in_cluster = cute::block_rank_in_cluster();
bool is_first_cta_in_cluster = IsSchedDynamicPersistent ? (cta_rank_in_cluster == 0) : true;
bool is_first_cta_in_cluster = cta_rank_in_cluster == 0;
int cta_coord_v = cta_rank_in_cluster % size<0>(typename TiledMma::AtomThrID{});
bool is_mma_leader_cta = cta_coord_v == 0;
constexpr bool has_mma_peer_cta = size(AtomThrShapeMNK{}) == 2;
@@ -508,12 +580,57 @@ public:
bool is_epi_load_needed = collective_epilogue.is_producer_load_needed();
IsParticipant is_participant = {
(warp_category == WarpCategory::MMA), // mma
(warp_category == WarpCategory::Sched) && is_first_cta_in_cluster, // sched
(warp_category == WarpCategory::Sched) && (IsSchedDynamicPersistent ? is_first_cta_in_cluster : true), // sched
(warp_category == WarpCategory::MainloopLoad), // main_load
(warp_category == WarpCategory::EpilogueLoad) && is_epi_load_needed, // epi_load
(warp_category == WarpCategory::Epilogue) // epilogue
(warp_category == WarpCategory::Epilogue), // epilogue
(warp_category == WarpCategory::TensorMapUpdater) && IsTensorMapUpdateAsync // tensor_map_updater
};
int32_t sm_id = static_cast<int32_t>(cutlass::arch::SmId());
if constexpr (IsGroupedGemmKernel) {
// In case user wants to engage less SMs than available on device
sm_id = blockIdx.x + (blockIdx.y * gridDim.x);
}
auto tensormaps_init_main_load = [&] () {
if constexpr (IsTensorMapUpdateAsync) {
return collective_mainloop.template tensormaps_init<IsTensorMapUpdateAsync>(
params.mainloop,
shared_storage.tensormaps[0].mainloop,
params.hw_info.sm_count,
sm_id
);
}
else {
return nullptr;
}
};
auto tensormaps_init_epi_load = [&] () {
if constexpr (IsTensorMapUpdateAsync) {
return collective_epilogue.template tensormaps_init<true /*IsLoad*/, IsTensorMapUpdateAsync>(
params.epilogue,
shared_storage.tensormaps[0].epilogue,
params.hw_info.sm_count,
sm_id
);
}
else {
return nullptr;
}
};
decltype(tensormaps_init_main_load()) pre_init_main_load_tensormaps;
decltype(tensormaps_init_epi_load()) pre_init_epi_load_tensormaps;
if (is_participant.main_load) {
pre_init_main_load_tensormaps = tensormaps_init_main_load();
}
if (is_participant.epi_load) {
pre_init_epi_load_tensormaps = tensormaps_init_epi_load();
}
// Mainloop Load pipeline
typename MainloopPipeline::Params mainloop_pipeline_params;
if (WarpCategory::MainloopLoad == warp_category) {
@@ -581,9 +698,14 @@ public:
clc_pipeline_params.transaction_bytes = CLCResponseSize;
}
else {
clc_pipeline_params.consumer_arv_count = NumMainloopLoadThreads + NumEpilogueThreads + NumMMAThreads;
if (is_epi_load_needed) {
clc_pipeline_params.consumer_arv_count += NumEpilogueLoadThreads;
if constexpr (IsTensorMapUpdateAsync) {
clc_pipeline_params.consumer_arv_count = NumThreadsPerWarp;
}
else {
clc_pipeline_params.consumer_arv_count = NumMainloopLoadThreads + NumEpilogueThreads + NumMMAThreads;
if (is_epi_load_needed) {
clc_pipeline_params.consumer_arv_count += NumEpilogueLoadThreads;
}
}
}
// Now declare the pipeline outside the if constexpr
@@ -596,6 +718,32 @@ public:
}
}();
auto tensor_map_ready_pipeline = [&] () {
if constexpr (IsGroupedGemmKernel) {
// TMA update ready pipeline
typename TensorMapReadyPipeline::Params tensor_map_ready_pipeline_params;
if (WarpCategory::TensorMapUpdater == warp_category) {
tensor_map_ready_pipeline_params.role = TensorMapReadyPipeline::ThreadCategory::Producer;
}
else {
tensor_map_ready_pipeline_params.role = TensorMapReadyPipeline::ThreadCategory::Consumer;
}
tensor_map_ready_pipeline_params.initializing_warp = 8;
tensor_map_ready_pipeline_params.producer_arv_count = NumThreadsPerWarp;
tensor_map_ready_pipeline_params.consumer_arv_count = NumMainloopLoadThreads + NumEpilogueThreads + NumMMAThreads;
if (is_epi_load_needed) {
tensor_map_ready_pipeline_params.consumer_arv_count += NumEpilogueLoadThreads;
}
return TensorMapReadyPipeline(shared_storage.pipelines.tensor_map_ready, tensor_map_ready_pipeline_params);
}
else {
return clc_pipeline;
}
}();
// Mainloop-Epilogue pipeline
typename AccumulatorPipeline::Params accumulator_pipeline_params;
if (WarpCategory::MMA == warp_category) {
@@ -672,19 +820,66 @@ public:
CLCPipelineState clc_pipe_consumer_state;
CLCPipelineState clc_pipe_producer_state = cutlass::make_producer_start_state<CLCPipeline>();
TensorMapReadyPipelineState tensor_map_ready_pipe_consumer_state;
TensorMapReadyPipelineState tensor_map_ready_pipe_producer_state = cutlass::make_producer_start_state<TensorMapReadyPipeline>();
AccumulatorPipelineState accumulator_pipe_consumer_state;
AccumulatorPipelineState accumulator_pipe_producer_state = cutlass::make_producer_start_state<AccumulatorPipeline>();
dim3 block_id_in_cluster = cute::block_id_in_cluster();
int32_t sm_id = static_cast<int32_t>(cutlass::arch::SmId());
// Calculate mask after cluster barrier arrival
mainloop_pipeline.init_masks(cluster_shape, block_id_in_cluster);
accumulator_pipeline.init_masks(cluster_shape, block_id_in_cluster);
// TileID scheduler
TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster);
typename TileScheduler::WorkTileInfo work_tile_info = scheduler.initial_work_tile_info(cluster_shape);
TileScheduler scheduler(
(!IsTensorMapUpdateAsync || is_participant.sched || is_participant.tensor_map_updater)
? &shared_storage.clc_response[0][0]
: &shared_storage.clc_response[1][0],
params.scheduler,
block_id_in_cluster
);
auto work_tile_info = [&] () {
if constexpr (!IsSchedDynamicPersistent) {
// Ensure that the prefetched kernel does not touch
// unflushed global memory prior to this instruction.
// For the static grouped scheduler, the problem shapes
// might be produced by a previous kernel in global memory.
cutlass::arch::wait_on_dependent_grids();
}
if constexpr (IsTensorMapUpdateAsync) {
return scheduler.initial_work_tile_info(cluster_shape, [] (typename TileScheduler::CLCResponse response) {
CLCResponseWithAdditionalInformation response_with_additional_info = response;
response_with_additional_info.TMA_stage = 0;
response_with_additional_info.batch_changed = 1;
return response_with_additional_info;
});
}
else {
return scheduler.initial_work_tile_info(cluster_shape);
}
} ();
auto get_tma_desc_offset = [] ([[maybe_unused]] const auto& tile_info) {
if constexpr (IsTensorMapUpdateAsync) {
return tile_info.TMA_stage;
}
else {
return 0;
}
};
auto get_tensormap = [] (auto& tensormaps, [[maybe_unused]] auto tma_desc_offset) {
if constexpr (IsTensorMapUpdateAsync) {
return tensormaps[tma_desc_offset];
}
else {
return tensormaps;
}
};
auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
//
@@ -698,44 +893,67 @@ public:
// When problem shapes are only on device, the grid launched may be larger than the total number of blocks across groups
return;
}
// In case user wants to engage less SMs than available on device
sm_id = blockIdx.x + (blockIdx.y * gridDim.x);
}
// Optionally append 1s until problem shape is rank-4 in case it is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(problem_shape.get_problem_shape(work_tile_info.L_idx), 1);
if (is_participant.main_load) {
auto load_inputs = collective_mainloop.load_init(
problem_shape_MNKL, params.mainloop,
shared_storage.tensors.mainloop,
shared_storage.tensormaps.mainloop,
params.hw_info.sm_count, sm_id, work_tile_info.L_idx);
// Ensure that the prefetched kernel does not touch
// unflushed global memory prior to this instruction
cutlass::arch::wait_on_dependent_grids();
auto load_inputs = collective_mainloop.load_init<IsTensorMapUpdateAsync>(
problem_shape_MNKL, params.mainloop,
shared_storage.tensors.mainloop,
shared_storage.tensormaps[0].mainloop,
params.hw_info.sm_count, sm_id, problem_shape.groups(), work_tile_info.L_idx);
bool do_load_order_arrive = is_epi_load_needed;
Tensor gA_mkl = get<0>(load_inputs);
// Fetch a copy of tensormaps for the CTA from Params
auto input_tensormaps = get<rank(load_inputs) - 1>(load_inputs);
auto input_tensormaps = [&] ([[maybe_unused]] auto inputs) {
if constexpr (IsTensorMapUpdateAsync) {
return pre_init_main_load_tensormaps;
}
else {
static constexpr size_t idx = rank(inputs) - 1;
return get<idx>(inputs);
}
} (load_inputs);
if constexpr (IsTensorMapUpdateAsync) {
// Register reconfiguration
arch::warpgroup_reg_dealloc<GenericRegisterRequirement>();
}
auto pad_inputs = [] (auto& inputs, [[maybe_unused]] auto tensormaps) {
if constexpr (IsTensorMapUpdateAsync) {
return cute::tuple_cat(inputs, cute::make_tuple(tensormaps));
}
else {
return inputs;
}
};
// Initial batch's tensor address update
// Even the first tile for a CTA can be from any of the batches.
// And during initialization of the first TMA descriptor on host, we don't initialize to the first batch due to that args value being device-only.
bool is_first_iteration = true;
bool did_batch_change = true;
bool requires_clc_query = true;
do {
auto tma_desc_offset = get_tma_desc_offset(work_tile_info);
int32_t curr_batch = idx2crd(work_tile_info.L_idx, shape<4>(gA_mkl)); // Usually just returns work_tile_info.L_idx;
if constexpr (IsGroupedGemmKernel) {
problem_shape_MNKL = append<4>(problem_shape.get_problem_shape(curr_batch), 1);
}
if (did_batch_change) {
if (IsTensorMapUpdateAsync ? is_first_iteration : did_batch_change) {
collective_mainloop.tensormaps_perform_update(
shared_storage.tensormaps.mainloop,
shared_storage.tensormaps[0].mainloop,
params.mainloop,
input_tensormaps,
get_tensormap(input_tensormaps, tma_desc_offset),
problem_shape,
curr_batch
);
@@ -763,10 +981,11 @@ public:
params.mainloop,
mainloop_pipeline,
mainloop_pipe_producer_state,
load_inputs,
pad_inputs(load_inputs, get_tensormap(input_tensormaps, tma_desc_offset)),
cta_coord_mnk,
k_tile_iter, k_tile_prologue,
did_batch_change
IsTensorMapUpdateAsync ? is_first_iteration : did_batch_change, // did_batch_change
curr_batch
);
mainloop_pipe_producer_state = mainloop_producer_state_next;
@@ -779,10 +998,11 @@ public:
params.mainloop,
mainloop_pipeline,
mainloop_pipe_producer_state,
load_inputs,
pad_inputs(load_inputs, get_tensormap(input_tensormaps, tma_desc_offset)),
cta_coord_mnk,
k_tile_iter_next, k_tile_count - k_tile_prologue,
false /* did_batch_change - prologue loads handle tensormap acquire */
false, /* did_batch_change - prologue loads handle tensormap acquire */
curr_batch
);
mainloop_pipe_producer_state = mainloop_producer_state_next_;
@@ -791,16 +1011,17 @@ public:
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
work_tile_info,
clc_pipeline,
clc_pipe_consumer_state
tensor_map_ready_pipeline,
tensor_map_ready_pipe_consumer_state
);
work_tile_info = next_work_tile_info;
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
requires_clc_query = increment_pipe;
if (increment_pipe) {
++clc_pipe_consumer_state;
++tensor_map_ready_pipe_consumer_state;
}
// For subsequent tiles, check if batch changes and therefore, we need tensormap updates
is_first_iteration = false;
did_batch_change = curr_batch != idx2crd(work_tile_info.L_idx, shape<4>(gA_mkl));
} while (work_tile_info.is_valid());
collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state);
@@ -808,6 +1029,12 @@ public:
}
else if (is_participant.sched) {
if constexpr (IsTensorMapUpdateAsync) {
// Register reconfiguration
arch::warpgroup_reg_dealloc<GenericRegisterRequirement>();
}
// Grouped GEMM uses static tile scheduler
if constexpr (IsSchedDynamicPersistent) {
// Whether a new CLC query must be performed.
@@ -853,20 +1080,135 @@ public:
}
else {
cutlass::arch::wait_on_dependent_grids();
static_assert(IsTensorMapUpdateAsync || IsSchedDynamicPersistent, "We only support async tensor map update with static persistent scheduler");
auto update_tensor_map_stages = [&] (typename TileScheduler::CLCResponse next_work_tile_info_from_scheduler) {
if constexpr (IsTensorMapUpdateAsync) {
CLCResponseWithAdditionalInformation next_work_tile_info = next_work_tile_info_from_scheduler;
auto tensor_map_buffer_stage = work_tile_info.TMA_stage;
next_work_tile_info.batch_changed = work_tile_info.L_idx != next_work_tile_info.L_idx;
if (next_work_tile_info.batch_changed) {
++tensor_map_buffer_stage;
}
next_work_tile_info.TMA_stage = tensor_map_buffer_stage;
return next_work_tile_info;
}
};
do {
auto [next_work_tile_info, increment_pipe] = scheduler.advance_to_next_work(clc_pipeline, clc_pipe_producer_state);
auto [next_work_tile_info, increment_pipe] = scheduler.advance_to_next_work(clc_pipeline, clc_pipe_producer_state, 1, update_tensor_map_stages);
work_tile_info = next_work_tile_info;
if (increment_pipe) {
++clc_pipe_producer_state;
}
} while (work_tile_info.is_valid());
clc_pipeline.producer_tail(clc_pipe_producer_state);
// Push additional invalid work items for all tensormap updater threads
for (int i = 0; i < NumTensorMapUpdaterThreads / NumThreadsPerWarp;) {
auto [next_work_tile_info, increment_pipe] = scheduler.advance_to_next_work(clc_pipeline, clc_pipe_producer_state, 1, update_tensor_map_stages);
work_tile_info = next_work_tile_info;
if (increment_pipe) {
++clc_pipe_producer_state;
++i;
}
}
}
}
else if (is_participant.tensor_map_updater) {
if constexpr (IsTensorMapUpdateAsync) {
// Register reconfiguration
arch::warpgroup_reg_dealloc<GenericRegisterRequirement>();
}
// Ensure that the prefetched kernel does not touch
// unflushed global memory prior to this instruction
cutlass::arch::wait_on_dependent_grids();
if constexpr (IsTensorMapUpdateAsync) {
auto updater_id = canonical_warp_idx_sync() - static_cast<int>(WarpCategory::TensorMapUpdater);
clc_pipe_consumer_state += updater_id;
tensor_map_ready_pipe_producer_state += updater_id;
auto tensormaps_mainloop = collective_mainloop.tensormaps_init<IsTensorMapUpdateAsync>(
params.mainloop,shared_storage.tensormaps[updater_id+1].mainloop, params.hw_info.sm_count, sm_id);
auto tensormaps_epilogue_load = collective_epilogue.template tensormaps_init<true /*IsLoad*/, IsTensorMapUpdateAsync /*IsTmaAsyncUpdate*/>(
params.epilogue, shared_storage.tensormaps[updater_id+1].epilogue, params.hw_info.sm_count, sm_id);
auto tensormaps_epilogue_store = collective_epilogue.template tensormaps_init<false /*IsLoad*/, IsTensorMapUpdateAsync /*IsTmaAsyncUpdate*/>(
params.epilogue, shared_storage.tensormaps[updater_id+1].epilogue, params.hw_info.sm_count, sm_id);
auto update_tensor_map_and_increment_pipe_if_needed = [&] (auto &next_work_tile_info, auto &increment_pipe) {
auto next_batch = next_work_tile_info.L_idx;
auto did_batch_change = next_work_tile_info.batch_changed;
if (increment_pipe) {
tensor_map_ready_pipeline.producer_acquire(tensor_map_ready_pipe_producer_state);
if (next_work_tile_info.is_valid() && did_batch_change) {
auto tma_desc_offset = get_tma_desc_offset(next_work_tile_info);
collective_mainloop.template tensormaps_perform_update<false /*WaitForInflightTmaRequests*/>(
shared_storage.tensormaps[updater_id+1].mainloop,
params.mainloop,
tensormaps_mainloop[tma_desc_offset],
problem_shape,
next_batch
);
if (collective_epilogue.is_producer_load_needed()) {
collective_epilogue.template tensormaps_perform_update<true /*IsLoad*/, false /*WaitForInflightTmaRequests*/>(
shared_storage.tensormaps[updater_id+1].epilogue,
params.epilogue,
tensormaps_epilogue_load[tma_desc_offset],
problem_shape,
next_batch
);
}
collective_epilogue.template tensormaps_perform_update<false /*IsLoad*/, false /*WaitForInflightTmaRequests*/>(
shared_storage.tensormaps[updater_id+1].epilogue,
params.epilogue,
tensormaps_epilogue_store[tma_desc_offset],
problem_shape,
next_batch
);
collective_mainloop.tensormaps_fence_acquire(tensormaps_mainloop[tma_desc_offset]);
if (collective_epilogue.is_producer_load_needed()) {
collective_epilogue.template tensormaps_fence_acquire<true /* IsEpiLoad */>(tensormaps_epilogue_load[tma_desc_offset]);
}
collective_epilogue.template tensormaps_fence_acquire<false /* IsEpiLoad */>(tensormaps_epilogue_store[tma_desc_offset]);
}
if (cute::elect_one_sync()) {
cute::tma_desc_commit_group();
shared_storage.clc_response[1][tensor_map_ready_pipe_producer_state.index()] = next_work_tile_info;
cutlass::arch::fence_view_async_shared();
cute::tma_desc_wait_group();
}
// Signal the other warps that the TMA update is complete
tensor_map_ready_pipeline.producer_commit(tensor_map_ready_pipe_producer_state);
tensor_map_ready_pipe_producer_state += (NumTensorMapUpdaterThreads / NumThreadsPerWarp);
clc_pipe_consumer_state += (NumTensorMapUpdaterThreads / NumThreadsPerWarp);
}
};
do {
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
work_tile_info,
clc_pipeline,
clc_pipe_consumer_state
);
update_tensor_map_and_increment_pipe_if_needed(next_work_tile_info, increment_pipe);
work_tile_info = next_work_tile_info;
} while (work_tile_info.is_valid());
}
}
else if (is_participant.mma) {
// Tmem allocation sequence
tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);
__syncwarp();
@@ -875,19 +1217,13 @@ public:
collective_mainloop.set_tmem_offsets(tmem_storage, tmem_base_ptr);
auto mma_inputs = collective_mainloop.mma_init(tmem_storage, shared_storage.tensors.mainloop);
if constexpr (IsTensorMapUpdateAsync) {
// Register reconfiguration
arch::warpgroup_reg_dealloc<GenericRegisterRequirement>();
}
do {
// Fetch next work tile
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
work_tile_info,
clc_pipeline,
clc_pipe_consumer_state
);
if (increment_pipe) {
++clc_pipe_consumer_state;
}
if constexpr (IsGroupedGemmKernel) {
problem_shape_MNKL = append<4>(problem_shape.get_problem_shape(work_tile_info.L_idx), 1);
}
@@ -915,6 +1251,17 @@ public:
}
++accumulator_pipe_producer_state;
// Fetch next work tile
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
work_tile_info,
tensor_map_ready_pipeline,
tensor_map_ready_pipe_consumer_state
);
if (increment_pipe) {
++tensor_map_ready_pipe_consumer_state;
}
work_tile_info = next_work_tile_info;
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
} while (work_tile_info.is_valid());
@@ -949,6 +1296,12 @@ public:
}
else if (is_participant.epi_load) {
if constexpr (IsTensorMapUpdateAsync) {
// Register reconfiguration
arch::warpgroup_reg_dealloc<GenericRegisterRequirement>();
}
// Ensure that the prefetched kernel does not touch
// unflushed global memory prior to this instruction
cutlass::arch::wait_on_dependent_grids();
@@ -958,37 +1311,42 @@ public:
int current_wave = 0;
// Fetch a copy of tensormaps for the CTA from Params
auto epi_load_tensormap = get<0>(collective_epilogue.load_init(
params.epilogue, shared_storage.tensormaps.epilogue, params.hw_info.sm_count, sm_id));
auto epi_load_tensormap = [&] () {
if constexpr (IsTensorMapUpdateAsync) {
collective_epilogue.template load_init<IsTensorMapUpdateAsync>(
params.epilogue,
shared_storage.tensormaps[0].epilogue,
params.hw_info.sm_count,
sm_id
);
return pre_init_epi_load_tensormaps;
}
else {
return get<0>(collective_epilogue.template load_init<IsTensorMapUpdateAsync>(
params.epilogue, shared_storage.tensormaps[0].epilogue, params.hw_info.sm_count, sm_id));
}
} ();
// Initial batch's tensor address update
// Even the first tile for a CTA can be from any of the batches.
// And during initialization of the first TMA descriptor on host, we don't initialize to the first batch due to that args value being device-only.
bool is_first_iteration = true;
bool did_batch_change = true;
constexpr bool IsEpiLoad = true;
do {
int32_t curr_batch = work_tile_info.L_idx;
if (did_batch_change) {
collective_epilogue.template tensormaps_perform_update<IsEpiLoad>(
shared_storage.tensormaps.epilogue,
auto tma_desc_offset = get_tma_desc_offset(work_tile_info);
if (IsTensorMapUpdateAsync ? is_first_iteration : did_batch_change) {
collective_epilogue.template tensormaps_perform_update<IsEpiLoad, !IsTensorMapUpdateAsync>(
shared_storage.tensormaps[0].epilogue,
params.epilogue,
epi_load_tensormap,
get_tensormap(epi_load_tensormap, tma_desc_offset),
problem_shape,
curr_batch
);
}
bool compute_epilogue = TileScheduler::compute_epilogue(work_tile_info, params.scheduler);
// Get current work tile and fetch next work tile
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
work_tile_info,
clc_pipeline,
clc_pipe_consumer_state
);
work_tile_info = next_work_tile_info;
if (increment_pipe) {
++clc_pipe_consumer_state;
}
if (compute_epilogue) {
if (do_load_order_wait) {
@@ -1009,7 +1367,7 @@ public:
TileShape{},
TiledMma{},
shared_storage.tensors.epilogue,
cute::make_tuple(epi_load_tensormap, did_batch_change),
cute::make_tuple(get_tensormap(epi_load_tensormap, tma_desc_offset), IsTensorMapUpdateAsync ? is_first_iteration : did_batch_change),
reverse_epi_n
);
@@ -1017,9 +1375,21 @@ public:
}
current_wave++;
// Fetch the next work tile
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
work_tile_info,
tensor_map_ready_pipeline,
tensor_map_ready_pipe_consumer_state
);
work_tile_info = next_work_tile_info;
if (increment_pipe) {
++tensor_map_ready_pipe_consumer_state;
}
// Calculate the cta coordinates of the next work tile
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
// For subsequent tiles, check if batch changes and therefore, we need tensormap updates
is_first_iteration = false;
did_batch_change = curr_batch != work_tile_info.L_idx;
} while (work_tile_info.is_valid());
@@ -1035,6 +1405,11 @@ public:
}
else if (is_participant.epilogue) {
if constexpr (IsTensorMapUpdateAsync) {
// Register reconfiguration
arch::warpgroup_reg_alloc<AccumRegisterRequirement>();
}
// Wait for tmem allocate here
tmem_allocation_result_barrier.arrive_and_wait();
uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;
@@ -1043,20 +1418,43 @@ public:
auto warp_idx_in_epi = canonical_warp_idx_sync() - static_cast<int>(WarpCategory::Epilogue);
bool do_tail_store = false;
// Fetch a copy of tensormaps for the CTA from Params
auto epi_store_tensormap = get<0>(collective_epilogue.store_init(
params.epilogue, shared_storage.tensormaps.epilogue, params.hw_info.sm_count, sm_id));
auto epi_store_tensormap = [&] () {
if constexpr (IsTensorMapUpdateAsync) {
collective_epilogue.template store_init<IsTensorMapUpdateAsync /*IsTmaAsyncUpdate*/>(
params.epilogue,
shared_storage.tensormaps[0].epilogue,
params.hw_info.sm_count,
sm_id
);
return collective_epilogue.template tensormaps_init<false /*IsLoad*/, IsTensorMapUpdateAsync /*IsTmaAsyncUpdate*/>(
params.epilogue,
shared_storage.tensormaps[0].epilogue,
params.hw_info.sm_count,
sm_id,
warp_idx_in_epi == 0
);
}
else {
return get<0>(collective_epilogue.template store_init<IsTensorMapUpdateAsync>(
params.epilogue, shared_storage.tensormaps[0].epilogue, params.hw_info.sm_count, sm_id));
}
} ();
// Initial batch's tensor address update
// Even the first tile for a CTA can be from any of the batches.
// And during initialization of the first TMA descriptor on host, we don't initialize to the first batch due to that args value being device-only.
bool is_first_iteration = true;
bool did_batch_change = true;
constexpr bool IsEpiLoad = false;
do {
int32_t curr_batch = work_tile_info.L_idx;
if (did_batch_change && warp_idx_in_epi == 0) {
collective_epilogue.template tensormaps_perform_update<IsEpiLoad>(
shared_storage.tensormaps.epilogue,
auto tma_desc_offset = get_tma_desc_offset(work_tile_info);
if ((IsTensorMapUpdateAsync ? is_first_iteration : did_batch_change) && warp_idx_in_epi == 0) {
collective_epilogue.template tensormaps_perform_update<IsEpiLoad, !IsTensorMapUpdateAsync>(
shared_storage.tensormaps[0].epilogue,
params.epilogue,
epi_store_tensormap,
get_tensormap(epi_store_tensormap, tma_desc_offset),
problem_shape,
curr_batch
);
@@ -1064,12 +1462,12 @@ public:
// Fetch next work tile
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
work_tile_info,
clc_pipeline,
clc_pipe_consumer_state
tensor_map_ready_pipeline,
tensor_map_ready_pipe_consumer_state
);
if (increment_pipe) {
++clc_pipe_consumer_state;
++tensor_map_ready_pipe_consumer_state;
}
// Accumulator stage slice
@@ -1090,7 +1488,11 @@ public:
//
// Epilogue and write to gD
//
auto [load_state_next, store_state_next, acc_state_next] = collective_epilogue.template store<IsOverlappingAccum>(
auto [
load_state_next,
store_state_next,
acc_state_next
] = collective_epilogue.template store<IsOverlappingAccum>(
epi_load_pipeline,
epi_load_pipe_consumer_state,
epi_store_pipeline,
@@ -1104,7 +1506,7 @@ public:
TiledMma{},
accumulator,
shared_storage.tensors.epilogue,
cute::make_tuple(epi_store_tensormap, did_batch_change)
cute::make_tuple(get_tensormap(epi_store_tensormap, tma_desc_offset), IsTensorMapUpdateAsync ? is_first_iteration : did_batch_change)
);
epi_load_pipe_consumer_state = load_state_next;
epi_store_pipe_producer_state = store_state_next;
@@ -1114,6 +1516,7 @@ public:
work_tile_info = next_work_tile_info;
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
// For subsequent tiles, check if batch changes and therefore, we need tensormap updates
is_first_iteration = false;
did_batch_change = curr_batch != work_tile_info.L_idx;
} while (work_tile_info.is_valid());
@@ -1138,6 +1541,11 @@ public:
}
else {
if constexpr (IsTensorMapUpdateAsync) {
// Register reconfiguration
arch::warpgroup_reg_dealloc<GenericRegisterRequirement>();
}
}
}
};
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp
+1 -1
View File
@@ -222,7 +222,6 @@ public:
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Host facing host arguments
struct Arguments {
@@ -420,6 +419,7 @@ public:
using namespace cute;
using X = Underscore;
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
auto problem_shape = params.problem_shape;
// Account for multiple epilogue and transformation warps
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_mma_transform.hpp
+1 -1
View File
@@ -234,7 +234,6 @@ public:
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Host facing host arguments
struct Arguments {
@@ -493,6 +492,7 @@ public:
using namespace cute;
using X = Underscore;
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
auto problem_shape = params.problem_shape;
// Account for more than one epilogue warp
include/cutlass/gemm/kernel/sm100_gemm_cpasync_warpspecialized.hpp
+2 -1
View File
@@ -188,7 +188,6 @@ public:
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Host facing host arguments
struct Arguments {
@@ -371,6 +370,8 @@ public:
using namespace cute;
using X = Underscore;
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Separate out problem shape for convenience
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized.hpp
+1 -1
View File
@@ -210,7 +210,6 @@ public:
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Host facing host arguments
struct Arguments {
@@ -408,6 +407,7 @@ public:
using namespace cute;
using X = Underscore;
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Separate out problem shape for convenience
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp
+1 -1
View File
@@ -205,7 +205,6 @@ public:
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Device side arguments
struct Arguments {
@@ -396,6 +395,7 @@ public:
using namespace cute;
using X = Underscore;
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Separate out problem shape for convenience
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_mixed_input_transform.hpp
+1 -1
View File
@@ -202,7 +202,6 @@ public:
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Device side arguments
struct Arguments {
@@ -396,6 +395,7 @@ public:
using namespace cute;
using X = Underscore;
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Separate out problem shape for convenience
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_mma_transform.hpp
+1 -1
View File
@@ -216,7 +216,6 @@ public:
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Host facing host arguments
struct Arguments {
@@ -402,6 +401,7 @@ public:
using namespace cute;
using X = Underscore;
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Separate out problem shape for convenience
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
include/cutlass/gemm/kernel/sm100_sparse_gemm_tma_warpspecialized.hpp
+1 -1
View File
@@ -216,7 +216,6 @@ public:
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Host facing host arguments
struct Arguments {
@@ -456,6 +455,7 @@ public:
using namespace cute;
using X = Underscore;
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Separate out problem shape for convenience
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
include/cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp
+10 -9
View File
@@ -128,11 +128,11 @@ public:
scheduler_sm90(params.params_sm90_, clc_response_ptr) { }
// Returns the initial work tile info that will be computed over
template <typename ClusterShape>
template <typename ClusterShape, typename CallbackBeforeCommit = WorkTileInfo(*)(WorkTileInfo)>
CUTLASS_DEVICE
auto
initial_work_tile_info(ClusterShape cluster_shape) {
return scheduler_sm90.initial_work_tile_info(cluster_shape);
initial_work_tile_info(ClusterShape cluster_shape, CallbackBeforeCommit callback_before_commit = [] (WorkTileInfo info) { return info;}) {
return scheduler_sm90.initial_work_tile_info(cluster_shape, callback_before_commit);
}
template<class BlockShape, class ClusterShape>
@@ -189,15 +189,16 @@ public:
);
}
template <typename CLCPipeline, typename CLCPipelineState>
template <typename CLCPipeline, typename CLCPipelineState, typename CallbackBeforeCommit = WorkTileInfo(*)(WorkTileInfo)>
CUTLASS_DEVICE
auto
advance_to_next_work(
CLCPipeline& clc_pipeline,
CLCPipelineState clc_pipe_producer_state,
uint32_t advance_count = 1) {
uint32_t advance_count = 1,
CallbackBeforeCommit callback_before_commit = [] (WorkTileInfo info) { return info;}) {
return scheduler_sm90.advance_to_next_work(clc_pipeline, clc_pipe_producer_state, advance_count);
return scheduler_sm90.advance_to_next_work(clc_pipeline, clc_pipe_producer_state, advance_count, callback_before_commit);
}
//
@@ -304,11 +305,11 @@ public:
}
// Kernel helper function to get next CLC ID
template <class CLCPipeline, class CLCPipelineState>
template <class WorkTileWithCallbackInfo, class CLCPipeline, class CLCPipelineState>
CUTLASS_DEVICE
auto
fetch_next_work(
WorkTileInfo work_tile_info,
WorkTileWithCallbackInfo work_tile_info,
CLCPipeline& clc_pipeline,
CLCPipelineState clc_pipe_consumer_state) {
@@ -320,7 +321,7 @@ private:
// Methods
//
[[nodiscard]] CUTLASS_DEVICE
static CLCResponse
static auto
load_query_response(uint32_t smem_ptr) {
return UnderlyingScheduler::load_query_response(smem_ptr);
}
include/cutlass/gemm/kernel/sm100_tile_scheduler_stream_k.hpp
+40
View File
@@ -37,6 +37,8 @@
#include "cutlass/gemm/kernel/sm100_tile_scheduler.hpp"
#include "cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp"
#include "cutlass/gemm/kernel/tile_scheduler_params.h"
#include "cutlass/conv/detail.hpp"
////////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::kernel::detail {
@@ -177,6 +179,44 @@ public:
return params;
}
template <conv::Operator ConvOp, int NumSpatialDims, class TileShapeMNK, class AtomThrShape>
static Params
to_underlying_arguments(
cutlass::conv::ConvProblemShape<ConvOp, NumSpatialDims> problem_shape,
TileShapeMNK tile_shape_mnk,
AtomThrShape atom_thr_shape_mnk,
ClusterShape cluster_shape_mnk,
KernelHardwareInfo const& hw_info,
Arguments const& args,
void* workspace = nullptr
) {
auto problem_shape_mnkl = [&] () {
// Infer im2col linearization from ConvOp and TileShape
constexpr bool is_linearized_M = (ConvOp == conv::Operator::kFprop || ConvOp == conv::Operator::kDgrad)
&& cute::depth<0>(TileShapeMNK{}) == _0{};
constexpr bool is_linearized_K = ConvOp == conv::Operator::kWgrad && cute::depth<2>(TileShapeMNK{}) == _1{};
if constexpr (is_linearized_M || is_linearized_K) {
// transformation + im2col linearization
return cutlass::conv::detail::get_linearized_problem_shape_MNKL(problem_shape);
}
else {
// transformation
return cutlass::conv::detail::get_transformed_problem_shape_MNKL(problem_shape);
}
}();
return to_underlying_arguments(
problem_shape_mnkl,
tile_shape_mnk,
atom_thr_shape_mnk,
cluster_shape_mnk,
hw_info,
args,
workspace
);
}
static bool
can_implement(Arguments const& args) {
return UnderlyingStreamKScheduler::can_implement(args);
include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp
+1 -2
View File
@@ -233,7 +233,6 @@ public:
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Host facing host arguments
struct Arguments {
@@ -508,6 +507,7 @@ public:
using namespace cute;
using X = Underscore;
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
auto problem_shape = params.problem_shape;
// Account for more than one epilogue warp
@@ -551,7 +551,6 @@ public:
typename MainloopABPipeline::Params mainloop_ab_pipeline_params;
if (WarpCategory::MainloopABLoad == warp_category) {
mainloop_ab_pipeline_params.role = MainloopABPipeline::ThreadCategory::Producer;
// Initialize the barrier for TMA load prefetch
}
if (WarpCategory::MMA == warp_category) {
mainloop_ab_pipeline_params.role = MainloopABPipeline::ThreadCategory::Consumer;
include/cutlass/gemm/kernel/sm103_blockscaled_gemm_tma_warpspecialized.hpp
+2 -1
View File
@@ -212,7 +212,6 @@ public:
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Host facing host arguments
struct Arguments {
@@ -425,6 +424,7 @@ public:
using namespace cute;
using X = Underscore;
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Separate out problem shape for convenience
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
@@ -483,6 +483,7 @@ public:
if (WarpCategory::MainloopABLoad == warp_category) {
mainloop_ab_pipeline_params.role = MainloopABPipeline::ThreadCategory::Producer;
// Initialize the barrier for TMA load prefetch
}
if (WarpCategory::MMA == warp_category) {
mainloop_ab_pipeline_params.role = MainloopABPipeline::ThreadCategory::Consumer;
include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp
+33 -16
View File
@@ -135,7 +135,7 @@ public:
// Sink scheduler params as a member
Params scheduler_params;
SchedulerResponse *response_ptr_ = nullptr;
void *response_ptr_ = nullptr;
ProblemShape cached_problem_shapes_[2];
//
@@ -225,6 +225,8 @@ public:
for (int group = 0; group < groups; group++) {
auto ctas_along_m = cute::size(cute::ceil_div(cute::shape<0>(problem_shapes.get_host_problem_shape(group)), cute::shape<0>(cta_shape)));
auto ctas_along_n = cute::size(cute::ceil_div(cute::shape<1>(problem_shapes.get_host_problem_shape(group)), cute::shape<1>(cta_shape)));
if(ctas_along_m <= 0) ctas_along_m = 1;
if(ctas_along_n <= 0) ctas_along_n = 1;
auto problem_blocks_m = round_up(ctas_along_m, cute::get<0>(cluster_shape));
auto problem_blocks_n = round_up(ctas_along_n, cute::get<1>(cluster_shape));
total_ctas += problem_blocks_m * problem_blocks_n;
@@ -301,7 +303,7 @@ public:
int32_t log_swizzle_size,
RasterOrder raster_order) {
int32_t valid_tile = 1;
uint8_t valid_tile = 1;
// Use a warp to "speculatively" check if the work tile maps to the next 32 groups
int lane_idx = canonical_lane_idx();
@@ -329,7 +331,8 @@ public:
auto problem_blocks_n = round_up(ctas_along_n, (1 << log_swizzle_size) * cluster_shape.n());
group_info.problem_blocks_along_raster_order = raster_order == RasterOrder::AlongN ? problem_blocks_n : problem_blocks_m;
group_info.total_tiles = problem_blocks_m * problem_blocks_n;
} else {
}
else {
group_info.total_tiles = INT_MAX;
}
@@ -428,23 +431,31 @@ public:
scheduler_params.log_swizzle_size_,
scheduler_params.raster_order_);
}
template <typename TileSchedulerPipeline, typename TileSchedulerPipelineState>
template <typename TileSchedulerPipeline, typename TileSchedulerPipelineState, typename CallbackBeforeCommit = WorkTileInfo(*)(WorkTileInfo)>
CUTLASS_DEVICE
auto
advance_to_next_work(
TileSchedulerPipeline& scheduler_pipeline,
TileSchedulerPipelineState scheduler_pipe_producer_state,
uint32_t advance_count = 1) {
uint32_t advance_count = 1,
CallbackBeforeCommit callback_before_commit = [] (WorkTileInfo info) { return info;}) {
current_work_linear_idx_ += total_grid_size_ * uint64_t(advance_count);
auto work_tile = get_current_work_for_linear_idx(current_work_linear_idx_);
using WorkTileWithCallbackInfo = decltype(callback_before_commit(work_tile));
WorkTileWithCallbackInfo work_tile_with_callback_info = work_tile;
scheduler_pipeline.producer_acquire(scheduler_pipe_producer_state);
if (work_tile_with_callback_info.is_valid()) {
work_tile_with_callback_info = callback_before_commit(work_tile);
}
if (cute::elect_one_sync()) {
response_ptr_[scheduler_pipe_producer_state.index()] = work_tile;
reinterpret_cast<WorkTileWithCallbackInfo *>(response_ptr_)[scheduler_pipe_producer_state.index()] = work_tile_with_callback_info;
cutlass::arch::fence_view_async_shared();
scheduler_pipeline.producer_commit(scheduler_pipe_producer_state);
}
return cute::make_tuple(work_tile, true);
return cute::make_tuple(work_tile_with_callback_info, true);
}
// Returns whether the block assigned this work should compute the epilogue for the corresponding
@@ -555,31 +566,37 @@ public:
}
// Kernel helper function to get next work tile
template <typename TileSchedulerPipeline, typename TileSchedulerPipelineState>
template <typename WorkTileWithCallbackInfo, typename TileSchedulerPipeline, typename TileSchedulerPipelineState>
CUTLASS_DEVICE
auto
fetch_next_work(
WorkTileInfo work_tile_info,
WorkTileWithCallbackInfo work_tile_with_callback_info,
TileSchedulerPipeline& scheduler_pipeline,
TileSchedulerPipelineState scheduler_pipe_consumer_state) {
if (continue_current_work(work_tile_info)) {
return cute::make_tuple(work_tile_info, true);
if (continue_current_work(work_tile_with_callback_info)) {
return cute::make_tuple(work_tile_with_callback_info, true);
}
scheduler_pipeline.consumer_wait(scheduler_pipe_consumer_state);
auto work_tile = response_ptr_[scheduler_pipe_consumer_state.index()];
work_tile_with_callback_info = reinterpret_cast<WorkTileWithCallbackInfo *>(response_ptr_)[scheduler_pipe_consumer_state.index()];
cutlass::arch::fence_view_async_shared();
scheduler_pipeline.consumer_release(scheduler_pipe_consumer_state);
return cute::make_tuple(work_tile, true);
return cute::make_tuple(work_tile_with_callback_info, true);
}
// Returns the initial work tile info that will be computed over
template <class ClusterShape>
template <class ClusterShape, typename CallbackBeforeCommit = WorkTileInfo(*)(WorkTileInfo)>
CUTLASS_DEVICE
auto
initial_work_tile_info(ClusterShape) {
return get_current_work_for_linear_idx(current_work_linear_idx_);
initial_work_tile_info(ClusterShape, CallbackBeforeCommit callback_before_commit = [] (WorkTileInfo response) { return response;}) {
auto work_tile = get_current_work_for_linear_idx(current_work_linear_idx_);
using WorkTileWithCallbackInfo = decltype(callback_before_commit(work_tile));
WorkTileWithCallbackInfo work_tile_with_callback_info = work_tile;
if (work_tile_with_callback_info.is_valid()) {
work_tile_with_callback_info = callback_before_commit(work_tile);
}
return work_tile_with_callback_info;
}
};