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

v4.5 dev update. (#3153)

Browse Source
This commit is contained in:
Junkai-Wu
2026-04-07 12:16:05 -04:00
committed by GitHub
parent 418d38a5de
commit a221da7ccf
265 changed files with 4913 additions and 1478 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/conv/collective/builders/sm100_common.inl
+12
View File
@@ -186,6 +186,18 @@ sm100_make_tiled_mma() {
}
}
template <class ArchTag, int KernelSmemCarveout>
constexpr int sm100_reduced_smem_capacity_bytes() {
if constexpr (cute::is_same_v<ArchTag, arch::Sm100>) {
return cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
}
else if constexpr (cute::is_same_v<ArchTag, arch::Sm103>) {
return cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
}
else {
static_assert(cutlass::detail::dependent_false<ArchTag>, "Invalid ArchTag, only Sm10x are supported.");
}
}
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::conv::collective::detail
include/cutlass/conv/collective/builders/sm100_umma_builder.inl
+9 -4
View File
@@ -44,6 +44,7 @@ using namespace cute;
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ArchTag,
conv::Operator ConvOp,
class ElementA,
class GmemLayoutA,
@@ -58,7 +59,7 @@ template <
class KernelScheduleType
>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassTensorOp,
ConvOp,
ElementA,
@@ -73,6 +74,9 @@ struct CollectiveBuilder<
StageCountType,
KernelScheduleType,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
|| cute::is_same_v<ArchTag, arch::Sm103>
) &&
(cute::is_same_v<KernelScheduleType, KernelImplicitTmaWarpSpecialized1SmSm100> ||
cute::is_same_v<KernelScheduleType, KernelImplicitTmaWarpSpecialized2SmSm100> ||
cute::is_same_v<KernelScheduleType, KernelStridedDgradTmaWs1SmSm100> ||
@@ -191,12 +195,12 @@ private:
CLCResponseStorage +
TmemBasePtrsStorage);
// 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 = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
static constexpr int PipelineStages = detail::compute_stage_count_or_override<
Sm100ReducedSmemCapacityBytes, ElementAMma, ElementBMma, SmemTileShape>(StageCountType{});
ReducedSmemCapacityBytes, ElementAMma, ElementBMma, SmemTileShape>(StageCountType{});
constexpr static int NumSpatialDimensions = detail::gmem_layout_tags_to_spatial_dims<GmemLayoutA, GmemLayoutB>();
@@ -206,7 +210,8 @@ private:
NumSpatialDimensions,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>;
ClusterShape_MNK,
ArchTag>;
public:
using CollectiveOp = cutlass::conv::collective::CollectiveConv<
include/cutlass/conv/collective/sm100_implicit_gemm_umma_warpspecialized.hpp
+5 -2
View File
@@ -65,6 +65,7 @@ template <
int NumSpatialDims,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShapeMNKL_, // (MmaAtomShapeM, MmaAtomShapeN, TileK, optional: TileL)
class ElementA_,
@@ -79,7 +80,8 @@ struct CollectiveConv<
NumSpatialDims,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShapeMNKL_,
ElementA_,
ElementB_,
@@ -96,7 +98,8 @@ struct CollectiveConv<
NumSpatialDims,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = decltype(cute::take<0,3>(TileShapeMNKL_{})); // (MmaAtomShapeM, MmaAtomShapeN, TileK)
using ElementA = ElementA_;
using ElementB = ElementB_;
include/cutlass/conv/device/conv_universal_adapter.hpp
+2 -1
View File
@@ -351,7 +351,8 @@ public:
}
else {
if constexpr (ConvKernel::ArchTag::kMinComputeCapability == 100 ||
ConvKernel::ArchTag::kMinComputeCapability == 101) {
ConvKernel::ArchTag::kMinComputeCapability == 101
) {
launch_result = ClusterLauncher::launch_with_fallback_cluster(
grid,
cluster,
include/cutlass/conv/dispatch_policy.hpp
+3 -2
View File
@@ -116,14 +116,15 @@ template<
int NumSpatialDimensions_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = cute::Shape<cute::C<1>,cute::C<1>,cute::C<1>>
class ClusterShape_ = cute::Shape<cute::C<1>,cute::C<1>,cute::C<1>>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedImplicitGemm {
static constexpr int Stages = Stages_;
static constexpr int NumSpatialDimensions = NumSpatialDimensions_;
static constexpr Operator ConvOp = ConvOp_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelScheduleImplicitTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
static_assert(NumSpatialDimensions >= 1);
include/cutlass/conv/kernel/sm100_implicit_gemm_tma_warpspecialized.hpp
+1 -2
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 {
@@ -413,7 +412,7 @@ public:
// Kernel level shared memory storage
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// In a warp specialized kernel, collectives expose data movement and compute operations separately
CollectiveMainloop collective_mainloop(params.mainloop, cluster_shape, cta_rank_in_cluster);
CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
include/cutlass/detail/layout.hpp
+3
View File
@@ -310,6 +310,9 @@ constexpr bool is_tma_copy_engine() {
|| cute::is_base_of_v<cute::SM90_TMA_STORE_IM2COL, GmemTiledCopy>
|| cute::is_base_of_v<cute::SM100_TMA_2SM_LOAD, GmemTiledCopy>
|| cute::is_base_of_v<cute::SM100_TMA_2SM_LOAD_MULTICAST, GmemTiledCopy>
|| cute::is_base_of_v<cute::SM100_TMA_LOAD_2D_GATHER4, GmemTiledCopy>
|| cute::is_base_of_v<cute::SM100_TMA_LOAD_MULTICAST_2D_GATHER4, GmemTiledCopy>
|| cute::is_base_of_v<cute::SM100_TMA_STORE_2D_SCATTER4, GmemTiledCopy>
) {
return true;
}
include/cutlass/epilogue/collective/sm100_epilogue_array_planar_complex_tma_warpspecialized.hpp
+3 -1
View File
@@ -138,10 +138,12 @@ private:
constexpr static bool is_m_major_C = detail::is_m_major<StrideC>();
constexpr static bool is_m_major_D = detail::is_m_major<StrideD>();
static constexpr bool IsGroupedGemmKernel = !cute::is_same_v<InternalStrideD, StrideD>;
// 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 = NumMaxSchedulerPipelineStageCount + std::max(StagesC, (ReuseSmemC ? StagesC : StagesD)) + 2;
constexpr static uint32_t NumTmaDescriptorsPerSm = IsGroupedGemmKernel ? (NumMaxSchedulerPipelineStageCount + std::max(StagesC, (ReuseSmemC ? StagesC : StagesD)) + 2) : 1;
using SmemLayoutC = decltype(tile_to_shape(
SmemLayoutAtomC{},
include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp
+5 -5
View File
@@ -181,11 +181,6 @@ private:
// TMA store delay only benefits with loop unrolling
constexpr static bool DelayTmaStore = DelayTmaStore_ and UnrollEpiLoop;
// 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 = NumMaxSchedulerPipelineStageCount + std::max(StagesC, (ReuseSmemC ? StagesC : StagesD)) + 2;
struct CollectiveStorageWithC {
alignas(SmemAlignmentC) ArrayEngine<SmemElementC, cosize_v<SmemLayoutC>> smem_C;
alignas(SmemAlignmentD) ArrayEngine<SmemElementD, cosize_v<SmemLayoutD>> smem_D;
@@ -241,6 +236,11 @@ public:
static constexpr bool IsGroupedGemmKernel = !cute::is_same_v<InternalStrideD, StrideD>;
// 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 = IsGroupedGemmKernel ? (NumMaxSchedulerPipelineStageCount + std::max(StagesC, (ReuseSmemC ? StagesC : StagesD)) + 2) : 1;
// Host side epilogue arguments
struct Arguments {
typename FusionCallbacks::Arguments thread{};
include/cutlass/gemm/collective/builders/sm100_9xBF16_interleaved_complex_umma_builder.inl
+6 -5
View File
@@ -185,8 +185,7 @@ struct CollectiveBuilder<
// Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
static constexpr auto stage_info = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_fast_fp32<
ReducedSmemCapacityBytes, CtaTileShape_MNK, TiledMma, BuilderScheduleTag, UmmaMajorACompute,
NumComplexComponents, NumComputeMtxs
@@ -209,7 +208,8 @@ struct CollectiveBuilder<
ScalingFactor,
AccPromotionInterval,
ClusterShape_MNK,
AccumulatorCopyAtom>,
AccumulatorCopyAtom,
ArchTag>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedFastF32<
Load2TransformPipelineStageCount,
Transform2MmaPipelineStageCount,
@@ -219,8 +219,9 @@ struct CollectiveBuilder<
ScalingFactor,
AccPromotionInterval,
ClusterShape_MNK,
AccumulatorCopyAtom>
>;
AccumulatorCopyAtom,
ArchTag>
>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
TileShape_MNK,
include/cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl
+6 -5
View File
@@ -253,8 +253,7 @@ struct CollectiveBuilder<
// Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
static constexpr auto stage_info = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_fast_fp32<
ReducedSmemCapacityBytes, CtaTileShape_MNK, TiledMma, BuilderScheduleTag, UmmaMajorACompute,
/*Cmplx=*/ 1, /*Mtxs=*/ NumComputeMtxs
@@ -276,7 +275,8 @@ struct CollectiveBuilder<
ScalingFactor,
AccPromotionInterval,
ClusterShape_MNK,
AccumulatorCopyAtom>,
AccumulatorCopyAtom,
ArchTag>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedFastF32<
Load2TransformPipelineStageCount,
Transform2MmaPipelineStageCount,
@@ -286,8 +286,9 @@ struct CollectiveBuilder<
ScalingFactor,
AccPromotionInterval,
ClusterShape_MNK,
AccumulatorCopyAtom>
>;
AccumulatorCopyAtom,
ArchTag>
>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
TileShape_MNK,
include/cutlass/gemm/collective/builders/sm100_blockscaled_mixed_tma_cpasync_umma_builder.inl
+11 -5
View File
@@ -93,6 +93,7 @@ sm100_compute_stage_count_or_override_blockscaled_mixed_tma_cpasync(StageCountAu
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ArchTag,
class ElementPairA,
class GmemLayoutATag,
int AlignmentA,
@@ -106,7 +107,7 @@ template <
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassBlockScaledTensorOp,
ElementPairA,
GmemLayoutATag,
@@ -119,7 +120,11 @@ struct CollectiveBuilder<
ClusterShape_MNK, // Static cluster shape (_1, _1, _1)
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<cute::is_same_v<KernelMixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100, BuilderScheduleTag> >
cute::enable_if_t<
cute::is_same_v<KernelMixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100, BuilderScheduleTag> &&
(cute::is_same_v<ArchTag, arch::Sm100>
)
>
>
{
using ElementSFA = typename detail::blockscaled::blockscaled_type<BuilderScheduleTag, ElementPairA>::sf_type;
@@ -238,12 +243,12 @@ 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 = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<SmemShapeA_M, BlockTileB_N, SmemShapeA_K>;
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_blockscaled_mixed_tma_cpasync<
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 CollectiveOp = cutlass::gemm::collective::CollectiveMma<
@@ -251,7 +256,8 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>,
ClusterShape_MNK,
ArchTag>,
TileShape_MNK,
cute::tuple<ElementA, ElementSF>,
StridePairA,
include/cutlass/gemm/collective/builders/sm100_blockscaled_sparse_umma_builder.inl
+3 -3
View File
@@ -278,8 +278,7 @@ 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 ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, 0>();
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_blockscaled_sparse<
ReducedSmemCapacityBytes,
@@ -297,7 +296,8 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>;
ClusterShape_MNK,
ArchTag>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
include/cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl
+7 -5
View File
@@ -255,8 +255,7 @@ struct CollectiveBuilder<
>::KernelSmemCarveout;
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
@@ -271,20 +270,23 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>,
cutlass::gemm::MainloopSm100ArrayTmaUmmaWarpSpecializedBlockScaled<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>
>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedBlockScaled<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>
>;
include/cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl
+5 -4
View File
@@ -375,8 +375,7 @@ struct CollectiveBuilder<
>::KernelSmemCarveout;
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
using MainloopABPipelineStorage = typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage;
@@ -417,12 +416,14 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>,
ClusterShape_MNK,
ArchTag>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedBlockwiseScaling<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>>;
ClusterShape_MNK,
ArchTag>>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
include/cutlass/gemm/collective/builders/sm100_common.inl
+12
View File
@@ -1000,6 +1000,18 @@ struct TrivialBlockscaledMma<
TileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, Instr, BuilderScheduleTag>());
};
template <class ArchTag, int KernelSmemCarveout>
constexpr int sm100_reduced_smem_capacity_bytes() {
if constexpr (cute::is_same_v<ArchTag, arch::Sm100>) {
return cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
}
else if constexpr (cute::is_same_v<ArchTag, arch::Sm103>) {
return cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
}
else {
static_assert(cutlass::detail::dependent_false<ArchTag>, "Invalid ArchTag, only Sm10x are supported.");
}
}
} // namespace detail
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/builders/sm100_cpasync_umma_builder.inl
+3 -3
View File
@@ -142,8 +142,7 @@ struct CollectiveBuilder<
CLCResponseStorage);
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
using MainloopPipelineStorage = typename cutlass::PipelineUmmaConsumerAsync<1>::SharedStorage;
@@ -156,7 +155,8 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>,
ClusterShape_MNK,
ArchTag>,
TileShape_MNK,
ElementA,
cutlass::gemm::TagToStrideA_t<GmemLayoutATag>,
include/cutlass/gemm/collective/builders/sm100_interleaved_complex_umma_builder.inl
+5 -4
View File
@@ -210,8 +210,7 @@ struct CollectiveBuilder<
TensorMapStorage);
// Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
@@ -230,7 +229,8 @@ struct CollectiveBuilder<
AccumulatorPipelineStageCount,
TransformationStageCount,
ClusterShape_MNK,
AccumulatorCopyAtom
AccumulatorCopyAtom,
ArchTag
>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedInterleavedComplexTF32<
PipelineStages,
@@ -238,7 +238,8 @@ struct CollectiveBuilder<
AccumulatorPipelineStageCount,
TransformationStageCount,
ClusterShape_MNK,
AccumulatorCopyAtom
AccumulatorCopyAtom,
ArchTag
>
>;
include/cutlass/gemm/collective/builders/sm100_mixed_input_umma_builder.inl
+3 -3
View File
@@ -307,8 +307,7 @@ struct CollectiveBuilder<
TensorMapStorage);
// Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, 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<
@@ -325,7 +324,8 @@ struct CollectiveBuilder<
Transform2MmaPipelineStageCount,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
include/cutlass/gemm/collective/builders/sm100_mixed_tma_cpasync_umma_builder.inl
+9 -6
View File
@@ -41,6 +41,7 @@ namespace cutlass::gemm::collective {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ArchTag,
class ElementA,
class GmemLayoutATag,
int AlignmentA,
@@ -54,7 +55,7 @@ template <
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassTensorOp,
ElementA,
GmemLayoutATag,
@@ -67,7 +68,9 @@ struct CollectiveBuilder<
ClusterShape_MNK, // Static cluster shape (_1, _1, _1)
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<cute::is_same_v<KernelMixedTmaCpAsyncWarpSpecialized1SmSm100, BuilderScheduleTag> >
cute::enable_if_t<cute::is_same_v<KernelMixedTmaCpAsyncWarpSpecialized1SmSm100, BuilderScheduleTag> &&
(cute::is_same_v<ArchTag, arch::Sm100>
)>
>
{
static_assert(cute::is_static_v<TileShape_MNK>, "TileShape has to be static");
@@ -135,20 +138,20 @@ 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 = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<SmemShapeA_M, BlockTileB_N, SmemShapeA_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::MainloopSm100UmmaMixedTmaCpAsyncWarpSpecialized<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>,
ClusterShape_MNK,
ArchTag>,
TileShape_MNK,
ElementA,
cutlass::gemm::TagToStrideA_t<GmemLayoutATag>,
include/cutlass/gemm/collective/builders/sm100_planar_complex_umma_builder.inl
+5 -5
View File
@@ -132,9 +132,7 @@ struct CollectiveBuilder<
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
// Use complex type to calculate SMEM stage count
@@ -150,13 +148,15 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedPlanarComplex<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>
>;
include/cutlass/gemm/collective/builders/sm100_sparse_umma_builder.inl
+3 -3
View File
@@ -381,8 +381,7 @@ 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 ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, 0>();
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_sparse<
ReducedSmemCapacityBytes,
@@ -399,7 +398,8 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>,
ClusterShape_MNK,
ArchTag>,
TileShape_MNK,
ElementA,
LayoutPairAE,
include/cutlass/gemm/collective/builders/sm100_umma_builder.inl
+7 -6
View File
@@ -288,9 +288,7 @@ struct CollectiveBuilder<
IsArrayOfPointersGemm
>::KernelSmemCarveout;
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
@@ -307,20 +305,23 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>,
cutlass::gemm::MainloopSm100ArrayTmaUmmaWarpSpecialized<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>
>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecialized<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>
>;
include/cutlass/gemm/collective/builders/sm103_blockscaled_umma_builder.inl
+5 -5
View File
@@ -498,9 +498,7 @@ struct CollectiveBuilder<
TensorMapStorage +
TmaPrefetchStorage);
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, 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<
ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, SmemLayoutAtomSFA, SmemLayoutAtomSFB>(StageCountType{});
@@ -512,7 +510,8 @@ struct CollectiveBuilder<
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK,
PrefetchType
PrefetchType,
ArchTag
>,
cutlass::gemm::MainloopSm103TmaUmmaWarpSpecializedBlockScaled<
get<0>(PipelineStages),
@@ -520,7 +519,8 @@ struct CollectiveBuilder<
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK,
PrefetchType
PrefetchType,
ArchTag
>
>;
include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp
+10 -7
View File
@@ -62,6 +62,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementPairA_,
@@ -82,7 +83,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementPairA_,
StridePairA_,
@@ -108,7 +110,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
// Due to an MSVC bug, we can't use decltype(make_tiled_mma()) interface.
using TiledMMA_SF = TiledMMA<MMA_Atom<typename TiledMma::MMA_ScaleFactor>,
@@ -143,11 +146,6 @@ 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;
@@ -264,6 +262,11 @@ struct CollectiveMma<
static constexpr bool IsF8F6F4 = detail::is_sm100_mma_f8f6f4<TiledMma, ElementA, ElementB>();
static constexpr bool IsGroupedGemmKernel = !cute::is_same_v<InternalStrideA, StrideA>;
// 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 = IsGroupedGemmKernel ? (SchedulerPipelineStageCount + Stages + 2) : 1;
using TmaInternalElementA = cute::conditional_t<IsF8F6F4, ElementAMma, ElementA>;
using TmaInternalElementB = cute::conditional_t<IsF8F6F4, ElementBMma, ElementB>;
include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized_rcggemm.hpp
+5 -2
View File
@@ -64,6 +64,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementPairA_,
@@ -84,7 +85,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementPairA_,
StridePairA_,
@@ -110,7 +112,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
// Due to an MSVC bug, we can't use decltype(make_tiled_mma()) interface.
using TiledMMA_SF = TiledMMA<MMA_Atom<typename TiledMma::MMA_ScaleFactor>,
include/cutlass/gemm/collective/sm100_blockscaled_mma_mixed_tma_cpasync_warpspecialized.hpp
+5 -2
View File
@@ -63,6 +63,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementPairA_,
@@ -83,7 +84,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementPairA_,
StridePairA_,
@@ -111,7 +113,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
// TileShape refers to MmaTileShape to adapt for runtime cluster
using TileShape = TileShape_;
using TiledMma_SF = TiledMMA<MMA_Atom<typename TiledMma::MMA_ScaleFactor>,
include/cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp
+5 -2
View File
@@ -62,6 +62,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementPairA_,
@@ -82,7 +83,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementPairA_,
StridePairA_,
@@ -108,7 +110,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
using TiledMMA_SF = TiledMMA<MMA_Atom<typename TiledMma::MMA_ScaleFactor>,
Layout<Shape<_1,_1,_1>>,
include/cutlass/gemm/collective/sm100_blockscaled_sparse_mma_warpspecialized.hpp
+5 -2
View File
@@ -63,6 +63,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementPairA_,
@@ -83,7 +84,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementPairA_,
LayoutPairA_,
@@ -109,7 +111,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
using TiledMMA_SF = TiledMMA<MMA_Atom<typename TiledMma::MMA_ScaleFactor>,
Layout<Shape<_1,_1,_1>>,
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized.hpp
+10 -7
View File
@@ -63,6 +63,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -83,7 +84,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StrideA_,
@@ -109,7 +111,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
@@ -123,11 +126,6 @@ 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_;
@@ -227,6 +225,11 @@ struct CollectiveMma<
static constexpr bool IsGroupedGemmKernel = !cute::is_same_v<InternalStrideA, StrideA>;
// 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 = IsGroupedGemmKernel ? (SchedulerPipelineStageCount + Stages + 2) : 1;
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128, _0> {
cute::ArrayEngine<SmemAllocTypeA, cute::cosize_v<SmemLayoutA>> smem_A;
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_blockwise_scaling.hpp
+5 -2
View File
@@ -60,6 +60,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -80,7 +81,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StridePairA_,
@@ -106,7 +108,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp
+5 -2
View File
@@ -69,6 +69,7 @@ template <
int NumBandsToCompute_,
int ScalingFactor_,
int AccPromotionInterval_,
class ArchTag_,
class AccumulatorCopyAtom_,
class ClusterShape,
class TileShape_,
@@ -93,7 +94,8 @@ struct CollectiveMma<
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>,
AccumulatorCopyAtom_,
ArchTag_>,
TileShape_,
float,
StrideA_,
@@ -124,7 +126,8 @@ struct CollectiveMma<
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>;
AccumulatorCopyAtom_,
ArchTag_>;
using TileShape = TileShape_;
using TiledMma = TiledMma_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp
+5 -2
View File
@@ -68,6 +68,7 @@ template <
int NumBandsToCompute_,
int ScalingFactor_,
int AccPromotionInterval_,
class ArchTag_,
class AccumulatorCopyAtom_,
class ClusterShape,
class TileShape_,
@@ -92,7 +93,8 @@ struct CollectiveMma<
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>,
AccumulatorCopyAtom_,
ArchTag_>,
TileShape_,
complex<float>,
StrideA_,
@@ -160,7 +162,8 @@ public:
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>;
AccumulatorCopyAtom_,
ArchTag_>;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp
+5 -2
View File
@@ -65,6 +65,7 @@ template <
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
int TransformationPipelineStageCount_,
class ArchTag_,
class AccumulatorCopyAtom_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
@@ -86,7 +87,8 @@ struct CollectiveMma<
AccumulatorPipelineStageCount_,
TransformationPipelineStageCount_,
ClusterShape,
AccumulatorCopyAtom_>,
AccumulatorCopyAtom_,
ArchTag_>,
TileShape_,
complex<float>,
StrideA_,
@@ -144,7 +146,8 @@ public:
AccumulatorPipelineStageCount_,
TransformationPipelineStageCount_,
ClusterShape,
AccumulatorCopyAtom_>;
AccumulatorCopyAtom_,
ArchTag_>;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_planar_complex.hpp
+12 -7
View File
@@ -64,6 +64,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape,
class TileShape_, // Static cluster shape or dynamic (int, int, _1)
class ElementA_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
@@ -84,7 +85,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StrideA_,
@@ -116,7 +118,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
CUTE_STATIC_ASSERT_V(evenly_divides(TileShape{}, tile_shape(TiledMma{})),
@@ -128,15 +131,17 @@ 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_;
using InternalStrideA = cute::remove_pointer_t<StrideA>;
static constexpr bool IsGroupedGemmKernel = !cute::is_same_v<InternalStrideA, StrideA>;
// 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 = IsGroupedGemmKernel ? (SchedulerPipelineStageCount + Stages + 2) : 1;
using ElementB = ElementB_;
using ElementBMma = typename TiledMma::ValTypeB;
using StrideB = StrideB_;
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_rcggemm.hpp
+5 -2
View File
@@ -65,6 +65,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -85,7 +86,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StrideA_,
@@ -111,7 +113,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
include/cutlass/gemm/collective/sm100_mma_cpasync_warpspecialized.hpp
+5 -2
View File
@@ -60,6 +60,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -80,7 +81,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StrideA_,
@@ -107,7 +109,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
// TileShape refers to MmaTileShape to adapt for runtime cluster shape
using TileShape = TileShape_;
include/cutlass/gemm/collective/sm100_mma_mixed_tma_cpasync_warpspecialized.hpp
+5 -2
View File
@@ -62,6 +62,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -82,7 +83,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StrideA_,
@@ -111,7 +113,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
// TileShape refers to MmaTileShape to adapt for runtime cluster
using TileShape = TileShape_;
include/cutlass/gemm/collective/sm100_mma_warpspecialized.hpp
+5 -2
View File
@@ -62,6 +62,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -82,7 +83,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StrideA_,
@@ -108,7 +110,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
include/cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp
+5 -2
View File
@@ -62,6 +62,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -82,7 +83,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StridePairA_,
@@ -108,7 +110,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
using ElementA = ElementA_;
include/cutlass/gemm/collective/sm100_mma_warpspecialized_emulated.hpp
+5 -2
View File
@@ -82,6 +82,7 @@ template <
int NumBandsToCompute_,
int ScalingFactor_,
int AccPromotionInterval_,
class ArchTag_,
class AccumulatorCopyAtom_,
class ClusterShape,
class TileShape_,
@@ -106,7 +107,8 @@ struct CollectiveMma<
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>,
AccumulatorCopyAtom_,
ArchTag_>,
TileShape_,
float,
StrideA_,
@@ -137,7 +139,8 @@ struct CollectiveMma<
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>;
AccumulatorCopyAtom_,
ArchTag_>;
using TileShape = TileShape_;
using TiledMma = TiledMma_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp
+5 -2
View File
@@ -66,6 +66,7 @@ template <
int NumBandsToCompute_,
int ScalingFactor_,
int AccPromotionInterval_,
class ArchTag_,
class AccumulatorCopyAtom_,
class ClusterShape,
class TileShape_,
@@ -90,7 +91,8 @@ struct CollectiveMma<
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>,
AccumulatorCopyAtom_,
ArchTag_>,
TileShape_,
complex<float>,
StrideA_,
@@ -156,7 +158,8 @@ public:
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>;
AccumulatorCopyAtom_,
ArchTag_>;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp
+5 -2
View File
@@ -78,6 +78,7 @@ template <
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
int TransformationPipelineStageCount_,
class ArchTag_,
class AccumulatorCopyAtom_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
@@ -99,7 +100,8 @@ struct CollectiveMma<
AccumulatorPipelineStageCount_,
TransformationPipelineStageCount_,
ClusterShape,
AccumulatorCopyAtom_>,
AccumulatorCopyAtom_,
ArchTag_>,
TileShape_,
complex<float>,
StrideA_,
@@ -155,7 +157,8 @@ public:
AccumulatorPipelineStageCount_,
TransformationPipelineStageCount_,
ClusterShape,
AccumulatorCopyAtom_>;
AccumulatorCopyAtom_,
ArchTag_>;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
include/cutlass/gemm/collective/sm100_mma_warpspecialized_mixed_input.hpp
+5 -2
View File
@@ -65,6 +65,7 @@ template <
int Transform2MmaPipelineStageCount_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ArchTag_,
class ClusterShape,
class TileShape_,
class ElementAOptionalTuple_,
@@ -86,7 +87,8 @@ struct CollectiveMma<
Transform2MmaPipelineStageCount_,
SchedulerPipelineStageCount_,
AccumulatorPipelineStageCount_,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementAOptionalTuple_,
StridePairA_,
@@ -115,7 +117,8 @@ public:
Transform2MmaPipelineStageCount_,
SchedulerPipelineStageCount_,
AccumulatorPipelineStageCount_,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
using TiledMma = TiledMma_;
using KernelSchedule = typename DispatchPolicy::Schedule;
include/cutlass/gemm/collective/sm100_mma_warpspecialized_planar_complex.hpp
+5 -2
View File
@@ -70,6 +70,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape,
class TileShape_, // Static cluster shape or dynamic (int, int, _1)
class ElementA_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
@@ -90,7 +91,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StrideA_,
@@ -122,7 +124,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
CUTE_STATIC_ASSERT_V(evenly_divides(TileShape{}, tile_shape(TiledMma{})),
include/cutlass/gemm/collective/sm100_sparse_mma_warpspecialized.hpp
+5 -2
View File
@@ -61,6 +61,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -81,7 +82,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
LayoutPairAE_,
@@ -107,7 +109,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
include/cutlass/gemm/collective/sm103_blockscaled_mma_array_warpspecialized.hpp
+5 -2
View File
@@ -66,6 +66,7 @@ template <
int LoadSFPipelineStageCount,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, int)
cutlass::sm103::detail::KernelPrefetchType PrefetchType,
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
@@ -89,7 +90,8 @@ struct CollectiveMma<
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape,
PrefetchType>,
PrefetchType,
ArchTag_>,
TileShape_,
ElementPairA_,
StridePairA_,
@@ -117,7 +119,8 @@ struct CollectiveMma<
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape,
PrefetchType>;
PrefetchType,
ArchTag_>;
using TileShape = TileShape_;
// Due to an MSVC bug, we can't use decltype(make_tiled_mma()) interface.
include/cutlass/gemm/collective/sm103_blockscaled_mma_warpspecialized.hpp
+5 -2
View File
@@ -66,6 +66,7 @@ template <
int LoadSFPipelineStageCount,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, int)
cutlass::sm103::detail::KernelPrefetchType PrefetchType,
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
@@ -89,7 +90,8 @@ struct CollectiveMma<
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape,
PrefetchType>,
PrefetchType,
ArchTag_>,
TileShape_,
ElementPairA_,
StridePairA_,
@@ -117,7 +119,8 @@ struct CollectiveMma<
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape,
PrefetchType>;
PrefetchType,
ArchTag_>;
using TileShape = TileShape_;
// Due to an MSVC bug, we can't use decltype(make_tiled_mma()) interface.
include/cutlass/gemm/dispatch_policy.hpp
+66 -44
View File
@@ -976,12 +976,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100UmmaCpAsyncWarpSpecialized {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -989,12 +990,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100UmmaMixedTmaCpAsyncWarpSpecialized {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelMixedTmaCpAsyncWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
@@ -1003,12 +1005,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100UmmaMixedTmaCpAsyncWarpSpecializedBlockScaled {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelMixedTmaCpAsyncWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
@@ -1018,12 +1021,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecialized {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
@@ -1033,12 +1037,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedBlockwiseScaling {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelTmaWarpSpecializedMmaTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
@@ -1048,12 +1053,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100ArrayTmaUmmaWarpSpecializedBlockwiseScaling {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelPtrArrayTmaWarpSpecializedMmaTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
@@ -1063,12 +1069,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedBlockScaled {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelTmaWarpSpecializedBlockScaledSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1077,13 +1084,14 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedSparse {
constexpr static int Stages = Stages_;
constexpr static int MetadataS2TStages = 4;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelSparseTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1092,13 +1100,14 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedBlockScaledSparse {
constexpr static int Stages = Stages_;
constexpr static int MetadataS2TStages = 4;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelSparseTmaWarpSpecializedBlockScaledSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1132,7 +1141,8 @@ template<
class ClusterShape_ = Shape<_1,_1,_1>,
// The TMEM_LOAD atom to be used for loading local accumulator
// from TMEM to registers
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_32dp32b32x
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_32dp32b32x,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedFastF32 {
constexpr static int Load2TransformPipelineStageCount = Load2TransformPipelineStageCount_;
@@ -1143,7 +1153,7 @@ struct MainloopSm100TmaUmmaWarpSpecializedFastF32 {
constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::FastF32;
using ClusterShape = ClusterShape_;
using AccumulatorCopyAtom = AccumulatorCopyAtom_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelTmaWarpSpecializedInputTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
@@ -1163,7 +1173,8 @@ template<
// Transformation <-> MMA
int TransformationPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>,
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_16dp256b1x
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_16dp256b1x,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedInterleavedComplexTF32 {
constexpr static int ComputationPipelineStageCount = ComputationPipelineStageCount_;
@@ -1171,7 +1182,7 @@ struct MainloopSm100TmaUmmaWarpSpecializedInterleavedComplexTF32 {
constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::InterleavedComplexTF32;
using ClusterShape = ClusterShape_;
using AccumulatorCopyAtom = AccumulatorCopyAtom_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelTmaWarpSpecializedInputTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
@@ -1191,7 +1202,8 @@ template<
// Accmulator pipeline depth
int AccumulatorPipelineStageCount_,
// ClusterShape for the kernel
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedMixedInput {
constexpr static int Load2TransformPipelineStageCount = Load2TransformPipelineStageCount_;
@@ -1199,7 +1211,7 @@ struct MainloopSm100TmaUmmaWarpSpecializedMixedInput {
constexpr static int Transform2MmaPipelineStageCount = Transform2MmaPipelineStageCount_;
constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::MixedInput;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelTmaWarpSpecializedMixedInputTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
@@ -1214,12 +1226,13 @@ template<
int SchedulerPipelineStageCount_,
// Accmulator pipeline depth
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedPlanarComplex {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
@@ -1229,12 +1242,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100ArrayTmaUmmaWarpSpecialized {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = false;
using Schedule = KernelPtrArrayTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1244,12 +1258,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100RCGroupGemmTmaUmmaWarpSpecialized {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = false;
using Schedule = KernelPtrArrayTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1259,12 +1274,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100RCGroupGemmTmaUmmaWarpSpecializedBlockScaled {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelPtrArrayTmaWarpSpecializedBlockScaledSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1274,12 +1290,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100ArrayTmaUmmaWarpSpecializedBlockScaled {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelPtrArrayTmaWarpSpecializedBlockScaledSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1291,12 +1308,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100ArrayTmaUmmaWarpSpecializedPlanarComplex {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = false;
using Schedule = KernelPtrArrayTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1331,7 +1349,8 @@ template<
class ClusterShape_ = Shape<_1,_1,_1>,
// The TMEM_LOAD atom to be used for loading local accumulator
// from TMEM to registers
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_32dp32b32x
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_32dp32b32x,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100ArrayTmaUmmaWarpSpecializedFastF32 {
constexpr static int Load2TransformPipelineStageCount = Load2TransformPipelineStageCount_;
@@ -1342,7 +1361,7 @@ struct MainloopSm100ArrayTmaUmmaWarpSpecializedFastF32 {
constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::FastF32;
using ClusterShape = ClusterShape_;
using AccumulatorCopyAtom = AccumulatorCopyAtom_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelPtrArrayTmaWarpSpecializedInputTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
@@ -1362,7 +1381,8 @@ template<
// Transformation <-> MMA
int TransformationPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>,
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_16dp256b1x
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_16dp256b1x,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100ArrayTmaUmmaWarpSpecializedInterleavedComplexTF32 {
constexpr static int ComputationPipelineStageCount = ComputationPipelineStageCount_;
@@ -1370,7 +1390,7 @@ struct MainloopSm100ArrayTmaUmmaWarpSpecializedInterleavedComplexTF32 {
constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::InterleavedComplexTF32;
using ClusterShape = ClusterShape_;
using AccumulatorCopyAtom = AccumulatorCopyAtom_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelPtrArrayTmaWarpSpecializedInputTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
@@ -1385,13 +1405,14 @@ template<
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>,
cutlass::sm103::detail::KernelPrefetchType PrefetchType_ = cutlass::sm103::detail::KernelPrefetchType::TmaPrefetch
cutlass::sm103::detail::KernelPrefetchType PrefetchType_ = cutlass::sm103::detail::KernelPrefetchType::TmaPrefetch,
class ArchTag_ = arch::Sm103
>
struct MainloopSm103TmaUmmaWarpSpecializedBlockScaled {
constexpr static int LoadABPipelineStageCount = LoadABPipelineStageCount_;
constexpr static int LoadSFPipelineStageCount = LoadSFPipelineStageCount_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm103;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelTmaWarpSpecializedBlockScaledSm103<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
@@ -1407,13 +1428,14 @@ template<
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>,
cutlass::sm103::detail::KernelPrefetchType PrefetchType_ = cutlass::sm103::detail::KernelPrefetchType::TmaPrefetch
cutlass::sm103::detail::KernelPrefetchType PrefetchType_ = cutlass::sm103::detail::KernelPrefetchType::TmaPrefetch,
class ArchTag_ = arch::Sm103
>
struct MainloopSm103ArrayTmaUmmaWarpSpecializedBlockScaled {
constexpr static int LoadABPipelineStageCount = LoadABPipelineStageCount_;
constexpr static int LoadSFPipelineStageCount = LoadSFPipelineStageCount_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm103;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelPtrArrayTmaWarpSpecializedBlockScaledSm103<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp
+1 -1
View File
@@ -144,7 +144,7 @@ public:
using TileSchedulerParams = typename TileScheduler::Params;
static constexpr bool IsSchedDynamicPersistent = TileScheduler::IsDynamicPersistent;
static constexpr bool IsTensorMapUpdateAsync = not IsSchedDynamicPersistent;
static constexpr bool IsTensorMapUpdateAsync = IsGroupedGemmKernel && not IsSchedDynamicPersistent;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
static constexpr uint32_t MinTensorMapWorkspaceAlignment = 64;
include/cutlass/gemm/kernel/sm100_static_tile_scheduler.hpp
+11 -1
View File
@@ -44,8 +44,18 @@ public StaticPersistentTileScheduler<
public:
using BaseScheduler = StaticPersistentTileScheduler<StaticPersistentTileScheduler100>;
public:
using BaseScheduler::StaticPersistentTileScheduler;
using Params = PersistentTileSchedulerSm90Params;
// Explicit forwarding constructors replacing inheriting-constructor syntax
// (`using BaseScheduler::StaticPersistentTileScheduler;`) which newer CUDA
// host compilers reject in dependent-base contexts: the injected-class-name
// resolves to a type rather than a constructor name.
CUTLASS_HOST_DEVICE
StaticPersistentTileScheduler100() = default;
CUTLASS_DEVICE explicit
StaticPersistentTileScheduler100(Params const& params_)
: BaseScheduler(params_) {}
using RasterOrder = typename Params::RasterOrder;
using RasterOrderOptions = typename Params::RasterOrderOptions;
struct CLCResponse { uint32_t data[4] = {0}; };
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp
+4 -2
View File
@@ -242,8 +242,10 @@ public:
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid cluster count to " << max_active_clusters);
}
KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count, max_active_clusters};
KernelHardwareInfo hw_info = args.hw_info;
hw_info.sm_count = sm_count;
hw_info.max_active_clusters = max_active_clusters;
// Calculate workspace pointers
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
size_t workspace_offset = 0;
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp
+4 -2
View File
@@ -245,8 +245,10 @@ public:
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid cluster count to " << max_active_clusters);
}
KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count, max_active_clusters};
KernelHardwareInfo hw_info = args.hw_info;
hw_info.sm_count = sm_count;
hw_info.max_active_clusters = max_active_clusters;
// Calculate workspace pointers
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
size_t workspace_offset = 0;
include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_cooperative.hpp
+4 -2
View File
@@ -193,8 +193,10 @@ public:
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid cluster count to " << max_active_clusters);
}
KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count, max_active_clusters};
KernelHardwareInfo hw_info = args.hw_info;
hw_info.sm_count = sm_count;
hw_info.max_active_clusters = max_active_clusters;
TileSchedulerParams scheduler = TileScheduler::to_underlying_arguments(
problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info, args.scheduler, workspace);
include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_pingpong.hpp
+4 -2
View File
@@ -206,8 +206,10 @@ public:
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid cluster count to " << max_active_clusters);
}
KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count, max_active_clusters};
KernelHardwareInfo hw_info = args.hw_info;
hw_info.sm_count = sm_count;
hw_info.max_active_clusters = max_active_clusters;
TileSchedulerParams scheduler = TileScheduler::to_underlying_arguments(
problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info, args.scheduler, workspace);
include/cutlass/kernel_hardware_info.h
+19 -9
View File
@@ -81,21 +81,26 @@ struct KernelHardwareInfo {
}
// Query maximum number of active clusters that could co-exist on the target device
// based on kernel properties such as cluster dims and threadblock dims
// based on kernel properties such as cluster dims and threadblock dims.
// When a green context stream is provided, the occupancy query is scoped to the
// green context's SM partition, returning the max active clusters for that partition.
static inline int
query_device_max_active_clusters(
dim3 cluster_dims,
uint32_t threads_per_block,
void const* kernel_ptr) {
void const* kernel_ptr,
cudaStream_t stream = nullptr) {
int max_active_clusters = 0;
#if defined(CUTLASS_SM90_CLUSTER_LAUNCH_ENABLED)
ClusterLauncher::LaunchConfig cluster_launch_config = ClusterLauncher::make_cluster_launch_config(
cluster_dims /* minimum grid dim */, cluster_dims, {threads_per_block, 1, 1});
cluster_dims /* minimum grid dim */, cluster_dims, {threads_per_block, 1, 1},
0 /* smem_size */, stream /* green ctx stream or nullptr */);
// Given the kernel function and launch configuration, return the maximum number of clusters that could co-exist on the target device.
// When stream is a green context stream, this returns the max active clusters for that partition.
cudaError_t result = cudaOccupancyMaxActiveClusters(&max_active_clusters, kernel_ptr, &cluster_launch_config.launch_config);
if (result != cudaSuccess) {
CUTLASS_TRACE_HOST(
" cudaGetDevice() returned error "
" cudaOccupancyMaxActiveClusters() returned error "
<< cudaGetErrorString(result));
return 0;
}
@@ -108,26 +113,31 @@ struct KernelHardwareInfo {
#endif
}
// Simpler version of the above query function that fetches relevant information from the Kernel
// Simpler version of the above query function that fetches relevant information from the Kernel.
// When a green context stream is provided, the occupancy query is scoped to that partition.
template <typename Kernel>
static inline int
query_device_max_active_clusters() {
query_device_max_active_clusters(cudaStream_t stream = nullptr) {
dim3 cluster_dims(cute::size<0>(typename Kernel::ClusterShape{}),
cute::size<1>(typename Kernel::ClusterShape{}),
cute::size<2>(typename Kernel::ClusterShape{}));
uint32_t threads_per_block = Kernel::MaxThreadsPerBlock;
void const* kernel_ptr = (void*)(device_kernel<Kernel>);
return query_device_max_active_clusters(cluster_dims, threads_per_block, kernel_ptr);
return query_device_max_active_clusters(cluster_dims, threads_per_block, kernel_ptr, stream);
}
// Create a KernelHardwareInfo by querying device properties.
// When a green context stream is provided, max_active_clusters is queried
// against that stream's green context partition instead of the full device.
template <typename Kernel>
static inline KernelHardwareInfo
make_kernel_hardware_info(int const device_id = 0, int sm_count = 0, int max_active_clusters = 0) {
make_kernel_hardware_info(int const device_id = 0, int sm_count = 0, int max_active_clusters = 0,
cudaStream_t stream = nullptr) {
if (sm_count == 0) {
sm_count = query_device_multiprocessor_count(device_id);
}
if (max_active_clusters == 0) {
max_active_clusters = query_device_max_active_clusters<Kernel>();
max_active_clusters = query_device_max_active_clusters<Kernel>(stream);
}
return {device_id, sm_count, max_active_clusters};
}
include/cutlass/version.h
+2 -2
View File
@@ -35,8 +35,8 @@
#include <string>
#define CUTLASS_MAJOR 4
#define CUTLASS_MINOR 4
#define CUTLASS_PATCH 2
#define CUTLASS_MINOR 5
#define CUTLASS_PATCH 0
#ifdef CUTLASS_VERSIONS_GENERATED
#include "cutlass/version_extended.h"