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

Support for Group GEMM in CUTLASS Profiler for Geforce and Spark (#3092)

Browse Source 
Co-authored-by: dePaul Miller <23461061+depaulmillz@users.noreply.github.com>
This commit is contained in:
dePaul Miller
2026-03-06 20:36:29 -05:00
committed by GitHub
co-authored by dePaul Miller
parent e5fcd125a5
commit 73c59c055c
8 changed files with 88 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/gemm/collective/builders/sm120_blockscaled_mma_builder.inl
+30 -7
View File
@@ -186,13 +186,13 @@ struct CollectiveBuilder<
// Basic storage block for new Scaling Factor Layouts
using mnBasicBlockShape = Shape<_32,_4>;
using mnBasicBlockStride = Stride<_16,_4>;
using kBasicBlockShape = Shape<Int<SFVectorSize>, Int<MMA_NSF>>;
using kBasicBlockShape = Shape<Int<(int)SFVectorSize>, Int<MMA_NSF>>;
using kBasicBlockStride = Stride<_0, _1>;
using sSFA_shapeM = decltype(prepend(size<0>(TileShape_MNK{}) / Blk_MN{}, mnBasicBlockShape{}));
using sSF_strideMN = decltype(prepend( Blk_Elems{}, mnBasicBlockStride{}));
using sSFA_strideM = sSF_strideMN;
using sSF_shapeK = decltype(prepend(make_shape( Blk_SF{}/Int<MMA_NSF>{}, size<2>(TileShape_MNK{}) / Int<SFVectorSize>{} / Blk_SF{}), kBasicBlockShape{}));
using sSF_shapeK = decltype(prepend(make_shape( Blk_SF{}/Int<MMA_NSF>{}, size<2>(TileShape_MNK{}) / Int<(int)SFVectorSize>{} / Blk_SF{}), kBasicBlockShape{}));
using sSFA_strideK = decltype(prepend(make_stride( Int<MMA_NSF>{}, size<0>(TileShape_MNK{}) / Blk_MN{} * Blk_Elems{}), kBasicBlockStride{}));
using sSFA_shape = decltype(make_shape( sSFA_shapeM{}, sSF_shapeK{}));
@@ -209,11 +209,6 @@ struct CollectiveBuilder<
using SmemLayoutAtomsA = decltype(cute::make_tuple(SmemLayoutAtomA{}, SmemLayoutAtomSFA{}));
using SmemLayoutAtomsB = decltype(cute::make_tuple(SmemLayoutAtomB{}, SmemLayoutAtomSFB{}));
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_blockscaled<
detail::sm120_smem_capacity_bytes, SmemAllocTypeA, SmemAllocTypeB, TileShape_MNK, SmemLayoutAtomSFA, SmemLayoutAtomSFB>(StageCountType{});
static constexpr uint32_t SchedulerPipelineStageCount = 3;
using StrideA = cutlass::gemm::TagToStrideA_t<GmemLayoutATag>;
using StrideB = cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>;
using InternalStrideA = cute::remove_pointer_t<StrideA>;
@@ -232,6 +227,34 @@ struct CollectiveBuilder<
cute::is_base_of_v<KernelPtrArrayTmaWarpSpecializedPingpong, BuilderScheduleTag>,
"Invalid builder schedule tag for grouped GEMM");
static constexpr uint32_t SchedulerPipelineStageCount = 3;
static constexpr int CLCResponseSize = sizeof(typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm100<Shape<_1,_1,_1>,1>::CLCResponse{});
static constexpr auto SchedulerPipelineStorage = IsGroupedGemmKernel ? sizeof(cutlass::PipelineDetail::PipelineAsyncSharedStorage<8>)
: sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, Shape<_1,_1,_1>>::SharedStorage);
static constexpr auto CLCResponseStorage = IsGroupedGemmKernel ? 0 : (SchedulerPipelineStageCount *
CLCResponseSize);
static constexpr auto TensorMapStorage =
IsGroupedGemmKernel ? sizeof(cute::TmaDescriptor) * 2 /* We have two tensormaps smem */ :
0;
// TensorMapReady pipeline storage (specific to grouped/array kernels)
static constexpr auto TensorMapReadyPipelineStorage =
IsGroupedGemmKernel ? sizeof(typename cutlass::PipelineAsync<SchedulerPipelineStageCount>::SharedStorage) :
0;
static constexpr int ReducedSmemCapacityBytes = detail::sm120_smem_capacity_bytes -
SchedulerPipelineStorage -
TensorMapStorage -
TensorMapReadyPipelineStorage -
CLCResponseStorage;
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_blockscaled<
ReducedSmemCapacityBytes, SmemAllocTypeA, SmemAllocTypeB, TileShape_MNK, SmemLayoutAtomSFA, SmemLayoutAtomSFB>(StageCountType{});
using KernelSchedule = cute::conditional_t<IsGroupedGemmKernel,
// PtrArray
cute::conditional_t<IsCooperative,
include/cutlass/gemm/collective/builders/sm1xx_common.inl
+2
View File
@@ -760,6 +760,8 @@ select_instr() {
(SfVectorSize == 32 && cute::is_same_v<KernelScheduleAuto, BuilderScheduleTag>)
|| (SfVectorSize == 32 && cute::is_base_of_v<KernelScheduleBlockScaledGemmSm100, BuilderScheduleTag>)
|| (SfVectorSize == 32 && cute::is_base_of_v<KernelTmaWarpSpecializedPingpong, BuilderScheduleTag>)
|| (SfVectorSize == 32 && cute::is_base_of_v<KernelPtrArrayTmaWarpSpecializedCooperative, BuilderScheduleTag>)
|| (SfVectorSize == 32 && cute::is_base_of_v<KernelPtrArrayTmaWarpSpecializedPingpong, BuilderScheduleTag>)
|| (SfVectorSize == 32 && cute::is_base_of_v<KernelTmaWarpSpecializedCooperative, BuilderScheduleTag>)
|| (SfVectorSize == 32 && cute::is_base_of_v<KernelSchedulePtrArrayBlockScaledGemmSm100, BuilderScheduleTag>)
|| (SfVectorSize == 64 && cute::is_base_of_v<KernelScheduleBlockScaledSparseGemmSm100, BuilderScheduleTag>