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

Blockscaled Ragged Contiguous Grouped Gemm for MoEs (#2790)

Browse Source 
* Adding blockscaled ragged contiguous grouped gemm for MoEs

* cleaning up the example

* introduction to example improved

---------

Co-authored-by: Shreya Gaur <shgaur@dc2-container-xterm-012.prd.it.nvidia.com>
This commit is contained in:
Shreya Gaur
2025-11-26 17:16:49 -08:00
committed by GitHub
parent e67e63c331
commit 2052fd3885
6 changed files with 2220 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
include/cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl
View File

@@ -241,8 +241,9 @@ struct CollectiveBuilder<
static constexpr uint32_t AccumulatorPipelineStageCount = (MMA_N == 256) ? 1 : 2;
static constexpr bool IsArrayOfPointersGemm = cute::is_base_of_v<KernelSchedulePtrArrayBlockScaledGemmSm100, 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);
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,
@@ -265,13 +266,21 @@ struct CollectiveBuilder<
using DispatchPolicy =
cute::conditional_t<IsArrayOfPointersGemm,
cutlass::gemm::MainloopSm100ArrayTmaUmmaWarpSpecializedBlockScaled<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedBlockScaled<
cute::conditional_t<IsRCGroupGemm,
cutlass::gemm::MainloopSm100RCGroupGemmTmaUmmaWarpSpecializedBlockScaled<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
>,
cutlass::gemm::MainloopSm100ArrayTmaUmmaWarpSpecializedBlockScaled<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
>
>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedBlockScaled<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,

1
include/cutlass/gemm/collective/collective_mma.hpp
View File

@@ -56,6 +56,7 @@
#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_blockscaled_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"

1293
include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized_rcggemm.hpp Normal file
View File

File diff suppressed because it is too large Load Diff

15
include/cutlass/gemm/dispatch_policy.hpp
View File

@@ -1192,6 +1192,21 @@ struct MainloopSm100RCGroupGemmTmaUmmaWarpSpecialized {
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 MainloopSm100RCGroupGemmTmaUmmaWarpSpecializedBlockScaled {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelPtrArrayTmaWarpSpecializedBlockScaledSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
template<
int Stages_,