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

v4.3 tag release update. (#2789)

Browse Source
This commit is contained in:
Junkai-Wu
2025-11-20 20:49:44 -05:00
committed by GitHub
parent 406e078b29
commit 8cd5bef43a
225 changed files with 23229 additions and 2813 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/gemm/collective/fp8_accumulation.hpp
+50 -33
View File
@@ -56,7 +56,6 @@ struct GmmaFP8Accumulation {
static_assert(is_rmem<TensorAccum>::value , "Accumulator tensor must be rmem resident.");
private:
TensorAccum& accum_;
TensorAccum accum_temp_;
uint32_t accum_promotion_interval_; // defines the max num of executed MMAs after which accum should be promoted.
@@ -65,8 +64,10 @@ private:
uint32_t reset_accum_flag_; // accum needs to be zeroed or not.
// promote or `add` the partial accumulators to main accumulator (FADD).
template <class TensorAccumOrig>
CUTLASS_DEVICE
void promote_core() {
void promote_core(TensorAccumOrig &accum_) {
CUTE_STATIC_ASSERT_V(size(accum_) == size(accum_temp_));
warpgroup_wait<0>();
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(accum_); ++i) {
@@ -75,8 +76,10 @@ private:
}
// `multiply` scale the partial accumulators and `add` to main accumulator (FFMA).
template <class TensorAccumOrig>
CUTLASS_DEVICE
void scale_core(ElementAccumulator const &scale) {
void scale_core(TensorAccumOrig &accum_, ElementAccumulator const &scale) {
CUTE_STATIC_ASSERT_V(size(accum_) == size(accum_temp_));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(accum_); ++i) {
accum_(i) += accum_temp_(i) * scale;
@@ -84,16 +87,17 @@ private:
}
template <
class TensorAccumOrig,
class EngineScale,
class LayoutScale>
CUTLASS_DEVICE
void scale_core(const cute::Tensor<EngineScale, LayoutScale> &scale) {
void scale_core(TensorAccumOrig &accum_, const cute::Tensor<EngineScale, LayoutScale> &scale) {
using TensorScale = cute::Tensor<EngineScale, LayoutScale>;
static_assert(is_static<LayoutScale>::value, "Scale Layout should be static");
static_assert(is_rmem<TensorScale>::value , "Scale tensor must be rmem resident.");
static_assert(LayoutAccum{}.shape() == LayoutScale{}.shape(), "Accumulator and scale must have same shape.");
CUTE_STATIC_ASSERT_V(size(accum_) == size(accum_temp_));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(accum_); ++i) {
@@ -102,12 +106,13 @@ private:
}
template <
class TensorAccumOrig,
class EngineScaleA,
class LayoutScaleA,
class EngineScaleB,
class LayoutScaleB>
CUTLASS_DEVICE
void scale_core(const cute::Tensor<EngineScaleA, LayoutScaleA> &scaleA, const cute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
void scale_core(TensorAccumOrig &accum_, const cute::Tensor<EngineScaleA, LayoutScaleA> &scaleA, const cute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
using TensorScaleA = cute::Tensor<EngineScaleA, LayoutScaleA>;
using TensorScaleB = cute::Tensor<EngineScaleB, LayoutScaleB>;
@@ -116,8 +121,10 @@ private:
static_assert(is_rmem<TensorScaleA>::value, "ScaleA tensor must be rmem resident.");
static_assert(is_rmem<TensorScaleB>::value, "ScaleB tensor must be rmem resident.");
static_assert(LayoutAccum{}.shape() == LayoutScaleA{}.shape(), "Accumulator and scaleA must have same shape.");
static_assert(LayoutAccum{}.shape() == LayoutScaleB{}.shape(), "Accumulator and scaleB must have same shape.");
CUTE_STATIC_ASSERT_V(size(accum_) == size(accum_temp_));
CUTE_STATIC_ASSERT_V(size(accum_) == size(scaleA));
CUTE_STATIC_ASSERT_V(size(accum_) == size(scaleB));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(accum_); ++i) {
@@ -128,16 +135,15 @@ private:
public:
CUTLASS_DEVICE
GmmaFP8Accumulation(
TensorAccum &accum,
TensorAccum &accum_temp,
uint32_t accum_promotion_interval,
uint32_t mma_count_per_mainloop_iteration)
: accum_(accum),
: accum_temp_(accum_temp),
accum_promotion_interval_(accum_promotion_interval),
mma_count_per_mainloop_iteration_(mma_count_per_mainloop_iteration),
mma_count_(0),
reset_accum_flag_(0)
{
accum_temp_ = cute::make_fragment_like(accum);
}
//
@@ -160,21 +166,23 @@ public:
//
/// promote (add) the results from the MMA accumulators to main accumulator if needed.
template <class TensorAccumOrig>
CUTLASS_DEVICE
void promote_if_needed() {
void promote_if_needed(TensorAccumOrig &accum_) {
mma_count_ += mma_count_per_mainloop_iteration_;
reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
if (reset_accum_flag_) {
promote_core();
promote_core(accum_);
mma_count_ = 0;
}
}
/// promote (add) the residue results from the MMA accumulators to main accumulator if needed.
template <class TensorAccumOrig>
CUTLASS_DEVICE
void promote_residue_if_needed() {
void promote_residue_if_needed(TensorAccumOrig &accum_) {
if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) {
promote_core();
promote_core(accum_);
}
}
@@ -183,95 +191,104 @@ public:
//
/// scale (multiply_add) the results from the MMA accumulators to main accumulator if needed.
template <class TensorAccumOrig>
CUTLASS_DEVICE
void scale_if_needed(ElementAccumulator const &scale) {
void scale_if_needed(TensorAccumOrig &accum_, ElementAccumulator const &scale) {
mma_count_ += mma_count_per_mainloop_iteration_;
reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
if (reset_accum_flag_) {
scale_core(scale);
scale_core(accum_, scale);
mma_count_ = 0;
}
}
template <
class TensorAccumOrig,
class EngineScale,
class LayoutScale>
CUTLASS_DEVICE
void scale_if_needed(const cute::Tensor<EngineScale, LayoutScale> &scale) {
void scale_if_needed(TensorAccumOrig &accum_, const cute::Tensor<EngineScale, LayoutScale> &scale) {
mma_count_ += mma_count_per_mainloop_iteration_;
reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
if (reset_accum_flag_) {
scale_core(scale);
scale_core(accum_, scale);
mma_count_ = 0;
}
}
template <
class TensorAccumOrig,
class EngineScaleA,
class LayoutScaleA,
class EngineScaleB,
class LayoutScaleB>
CUTLASS_DEVICE
void scale_if_needed(const cute::Tensor<EngineScaleA, LayoutScaleA> &scaleA, const cute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
void scale_if_needed(TensorAccumOrig &accum_, const cute::Tensor<EngineScaleA, LayoutScaleA> &scaleA, const cute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
mma_count_ += mma_count_per_mainloop_iteration_;
reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
if (reset_accum_flag_) {
scale_core(scaleA, scaleB);
scale_core(accum_, scaleA, scaleB);
mma_count_ = 0;
}
}
/// scale (multiply_add) the results from the MMA accumulators to main accumulator without checking the counter.
template <class TensorAccumOrig>
CUTLASS_DEVICE
void scale(ElementAccumulator const &scale) {
scale_core(scale);
void scale(TensorAccumOrig &accum_, ElementAccumulator const &scale) {
scale_core(accum_, scale);
}
template <
class TensorAccumOrig,
class EngineScale,
class LayoutScale>
CUTLASS_DEVICE
void scale(const cute::Tensor<EngineScale, LayoutScale> &scale) {
scale_core(scale);
void scale(TensorAccumOrig &accum_, const cute::Tensor<EngineScale, LayoutScale> &scale) {
scale_core(accum_, scale);
}
template <
class TensorAccumOrig,
class EngineScaleA,
class LayoutScaleA,
class EngineScaleB,
class LayoutScaleB>
CUTLASS_DEVICE
void scale(const cute::Tensor<EngineScaleA, LayoutScaleA> &scaleA, const cute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
scale_core(scaleA, scaleB);
void scale(TensorAccumOrig &accum_, const cute::Tensor<EngineScaleA, LayoutScaleA> &scaleA, const cute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
scale_core(accum_, scaleA, scaleB);
}
/// scale (multiply_add) the residue results from the MMA accumulators to main accumulator if needed.
template <class TensorAccumOrig>
CUTLASS_DEVICE
void scale_residue_if_needed(ElementAccumulator const &scale) {
void scale_residue_if_needed(TensorAccumOrig &accum_, ElementAccumulator const &scale) {
if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) {
scale_core(scale);
scale_core(accum_, scale);
}
}
template <
class TensorAccumOrig,
class EngineScale,
class LayoutScale>
CUTLASS_DEVICE
void scale_residue_if_needed(const cute::Tensor<EngineScale, LayoutScale> &scale) {
void scale_residue_if_needed(TensorAccumOrig &accum_, const cute::Tensor<EngineScale, LayoutScale> &scale) {
if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) {
scale_core(scale);
scale_core(accum_, scale);
}
}
template <
class TensorAccumOrig,
class EngineScaleA,
class LayoutScaleA,
class EngineScaleB,
class LayoutScaleB>
CUTLASS_DEVICE
void scale_residue_if_needed(const cute::Tensor<EngineScaleA, LayoutScaleA> &scaleA, const cute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
void scale_residue_if_needed(TensorAccumOrig &accum_, const cute::Tensor<EngineScaleA, LayoutScaleA> &scaleA, const cute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) {
scale_core(scaleA, scaleB);
scale_core(accum_, scaleA, scaleB);
}
}
};
include/cutlass/gemm/collective/sm100_blockscaled_mma_mixed_tma_cpasync_warpspecialized.hpp
+18 -29
View File
@@ -166,8 +166,8 @@ struct CollectiveMma<
using ElementBMma = typename TiledMma::ValTypeB;
using StrideB = remove_cvref_t<decltype(get<0>(StridePairB{}))>;
static constexpr bool IsRuntimeDataTypeA = cute::is_same_v<ElementA, cutlass::type_erased_dynamic_float8_t>;
static constexpr bool IsRuntimeDataTypeB = cute::is_same_v<ElementB, cutlass::type_erased_dynamic_float8_t>;
static constexpr bool IsRuntimeDataTypeA = cute::is_same_v<ElementA, cutlass::type_erased_dynamic_float8_t> or cute::is_same_v<ElementA, cutlass::type_erased_dynamic_float4_t>;
static constexpr bool IsRuntimeDataTypeB = cute::is_same_v<ElementB, cutlass::type_erased_dynamic_float8_t> or cute::is_same_v<ElementB, cutlass::type_erased_dynamic_float4_t>;
static_assert(IsRuntimeDataTypeA == IsRuntimeDataTypeB,
"ElementA and ElementB should be both runtime or both static.");
@@ -308,13 +308,9 @@ struct CollectiveMma<
// Host side kernel arguments
struct Arguments {
ArrayElementA const* ptr_A{nullptr};
StrideA dA{};
ArrayElementB const* ptr_B{nullptr};
StrideB dB{};
ElementSF const* ptr_SFA{nullptr};
LayoutSFA layout_SFA{};
ElementSF const* ptr_SFB{nullptr};
LayoutSFB layout_SFB{};
RuntimeDataTypeA runtime_data_type_a{};
RuntimeDataTypeB runtime_data_type_b{};
};
@@ -356,7 +352,6 @@ struct CollectiveMma<
TMA_SFB tma_load_sfb;
ArrayElementB const* ptr_B{nullptr};
StrideB dB{};
LayoutSFA layout_SFA;
LayoutSFB layout_SFB;
@@ -392,9 +387,15 @@ struct CollectiveMma<
auto ptr_A = recast_ptr<TmaInternalElementA>(args.ptr_A);
auto ptr_B = recast_ptr<ElementBMma>(args.ptr_B);
Tensor tensor_a = make_tensor(ptr_A, make_layout(make_shape(M,K,L), args.dA));
Tensor tensor_sfa = make_tensor(args.ptr_SFA, args.layout_SFA);
Tensor tensor_sfb = make_tensor(args.ptr_SFB, args.layout_SFB);
const auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL);
const auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL);
auto shape_a = make_shape(M, K, L);
auto stride_a = cutlass::make_internal_packed_stride(StrideA{}, shape_a);
Tensor tensor_a = make_tensor(ptr_A, make_layout(shape_a, stride_a));
Tensor tensor_sfa = make_tensor(args.ptr_SFA, layout_SFA);
Tensor tensor_sfb = make_tensor(args.ptr_SFB, layout_SFB);
auto cluster_layout_vmnk = tiled_divide(make_layout(ClusterShape{}), make_tile(typename TiledMma::AtomThrID{}));
auto cluster_layout_sfb_vmnk = tiled_divide(make_layout(ClusterShape{}), make_tile(typename TiledMma_SF::AtomThrID{}));
@@ -426,9 +427,8 @@ struct CollectiveMma<
tma_load_sfa,
tma_load_sfb,
args.ptr_B,
args.dB,
args.layout_SFA,
args.layout_SFB,
layout_SFA,
layout_SFB,
args.runtime_data_type_a,
args.runtime_data_type_b
};
@@ -458,19 +458,6 @@ struct CollectiveMma<
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for CpAsync.\n");
}
// Check for SFA SFB layout requirement
const auto layout_sfa_ref = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL);
const auto layout_sfb_ref = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL);
implementable = implementable && (layout_sfa_ref == args.layout_SFA);
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: layout_SFA mismatch, layout_SFA needs to be K-major\n");
}
implementable = implementable && (layout_sfb_ref == args.layout_SFB);
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: layout_SFB mismatch, layout_SFB needs to be K-major\n");
}
return implementable;
}
@@ -628,10 +615,14 @@ struct CollectiveMma<
// Separate out problem shape for convenience
auto [M,N,K,L] = problem_shape_MNKL;
// Setting the stride of B
auto shape_b = make_shape(N, K, L);
StrideB stride_b = cutlass::make_internal_packed_stride(StrideB{}, shape_b);
// convert to subptr iterator if necessary
auto ptr_B = recast_ptr<ElementBMma>(params.ptr_B);
// Represent the full tensors
Tensor mB_nkl = make_tensor(make_gmem_ptr(ptr_B), make_shape(N,K,L), params.dB); //(n,k,l)
Tensor mB_nkl = make_tensor(make_gmem_ptr(ptr_B), shape_b, stride_b); //(n,k,l)
// Partition for cpasync
Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
@@ -1032,8 +1023,6 @@ protected:
RuntimeDataTypeA runtime_data_type_a_{};
RuntimeDataTypeB runtime_data_type_b_{};
// ClusterShape cluster_shape_;
// uint32_t block_rank_in_cluster_;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_rcggemm.hpp
+10 -9
View File
@@ -50,6 +50,8 @@
#include "cute/algorithm/gemm.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
#include "cutlass/detail/collective/moe_stride_utils.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
@@ -257,9 +259,7 @@ struct CollectiveMma<
// 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{};
};
@@ -296,9 +296,7 @@ struct CollectiveMma<
RuntimeDataTypeB runtime_data_type_b;
cute::TmaDescriptor* tensormaps;
ArrayElementA const* ptr_A;
InternalStrideA dA;
ArrayElementB const** ptr_B;
StrideB dB;
};
CUTLASS_DEVICE
@@ -343,7 +341,9 @@ struct CollectiveMma<
init_M = get<0>(problem_shape_MNK);
init_K_A = get<2>(problem_shape_MNK);
InternalStrideA stride_a = args.dA;
auto shape_a = make_shape(init_M, init_K_A, problem_shapes.groups());
InternalStrideA stride_a = cutlass::make_internal_packed_stride(InternalStrideA{}, shape_a);
InternalStrideB stride_b = InternalStrideB{};
// Batches/Groups are managed by using appropriate pointers to input matrices.
@@ -397,9 +397,7 @@ struct CollectiveMma<
args.runtime_data_type_b,
reinterpret_cast<cute::TmaDescriptor*>(workspace),
args.ptr_A,
args.dA,
reinterpret_cast<ArrayElementB const**>(args.ptr_B),
args.dB
reinterpret_cast<ArrayElementB const**>(args.ptr_B)
};
}
@@ -812,7 +810,10 @@ struct CollectiveMma<
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]);
auto internal_shape_b = make_shape(static_cast<int>(N), static_cast<int>(K), 1);
InternalStrideB stride_b = cutlass::make_internal_packed_stride(InternalStrideB{}, internal_shape_b);
Tensor tensor_b = make_tensor(ptr_B, make_shape(N,K,Int<1>{}), stride_b);
cute::detail::fill_tma_gmem_shape_stride(*observed_tma_load_b_, tensor_b,
prob_shape_B, prob_stride_B);
include/cutlass/gemm/collective/sm100_mma_mixed_tma_cpasync_warpspecialized.hpp
+8 -14
View File
@@ -247,9 +247,7 @@ struct CollectiveMma<
// Host side kernel arguments
struct Arguments {
ArrayElementA const* ptr_A{nullptr};
StrideA dA{};
ArrayElementB const* ptr_B{nullptr};
StrideB dB{};
RuntimeDataTypeA runtime_data_type_a{};
RuntimeDataTypeB runtime_data_type_b{};
};
@@ -271,7 +269,6 @@ struct CollectiveMma<
TMA_A tma_load_a;
ArrayElementB const* ptr_B{nullptr};
StrideB dB{};
RuntimeDataTypeA runtime_data_type_a;
RuntimeDataTypeB runtime_data_type_b;
@@ -299,7 +296,9 @@ struct CollectiveMma<
auto ptr_A = recast_ptr<TmaInternalElementA>(args.ptr_A);
auto ptr_B = recast_ptr<ElementBMma>(args.ptr_B);
Tensor tensor_a = make_tensor(ptr_A, make_layout(make_shape(M,K,L), args.dA));
auto shape_a = make_shape(M, K, L);
StrideA stride_a = cutlass::make_internal_packed_stride(StrideA{}, shape_a);
Tensor tensor_a = make_tensor(ptr_A, make_layout(shape_a, stride_a));
auto cluster_layout_vmnk = tiled_divide(make_layout(ClusterShape{}), make_tile(typename TiledMma::AtomThrID{}));
@@ -314,7 +313,6 @@ struct CollectiveMma<
return {
tma_load_a,
args.ptr_B,
args.dB,
args.runtime_data_type_a,
args.runtime_data_type_b
};
@@ -443,7 +441,11 @@ struct CollectiveMma<
auto [M,N,K,L] = problem_shape_MNKL;
// Represent the full tensors
Tensor mB_nkl = make_tensor(make_gmem_ptr(params.ptr_B), make_shape(N,K,L), params.dB); //(n,k,l)
auto shape_b = make_shape(N, K, L);
StrideB stride_b = cutlass::make_internal_packed_stride(StrideB{}, shape_b);
Tensor mB_nkl = make_tensor(make_gmem_ptr(params.ptr_B), shape_b, stride_b); //(n,k,l)
// Partition for cpasync
Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
@@ -571,14 +573,6 @@ struct CollectiveMma<
ProblemShape_MNKL effective_shape
) {
// Unpack from load_inputs
// GTensorB tBgB_nkl = get<0>(load_inputs);
// CTensorB cgB_nk = get<1>(load_inputs);
// STensorB sB = get<2>(load_inputs);
// ProblemShape_MNKL problem_shape_MNKL = get<3>(load_inputs);
// TiledCopyB gmem_to_smem_b_tiled_copy = get<4>(load_inputs);
// ThreadCopyB thr_copy_b = get<5>(load_inputs);
auto [
tBgB_nkl, cgB_nk, sB,
// problem_shape_MNKL,
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8.hpp
+5 -5
View File
@@ -531,8 +531,8 @@ struct CollectiveMma<
int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
GmmaFP8Accumulation accumulation(accum, mainloop_params.mma_promotion_interval, size<2>(tCrA));
auto accm_temp = cute::make_fragment_like(accum);
GmmaFP8Accumulation accumulation(accm_temp, mainloop_params.mma_promotion_interval, size<2>(tCrA));
warpgroup_fence_operand(accumulation());
CUTLASS_PRAGMA_UNROLL
for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0; --k_tile_prologue)
@@ -556,7 +556,7 @@ struct CollectiveMma<
}
warpgroup_commit_batch();
accumulation.promote_if_needed();
accumulation.promote_if_needed(accum);
++smem_pipe_read;
}
@@ -597,7 +597,7 @@ struct CollectiveMma<
warpgroup_wait<K_PIPE_MMAS>();
warpgroup_fence_operand(accumulation());
accumulation.promote_if_needed();
accumulation.promote_if_needed(accum);
pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it
@@ -606,7 +606,7 @@ struct CollectiveMma<
++smem_pipe_release;
}
accumulation.promote_residue_if_needed();
accumulation.promote_residue_if_needed(accum);
warpgroup_fence_operand(accumulation());
}
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+176 -204
View File
@@ -212,6 +212,8 @@ struct CollectiveMma<
static_assert(cute::is_same_v<ElementAccumulator, ElementBlockScale>,
"ElementAccumulator and ElementBlockScale should be same datatype");
using NumSplitsM = cute::C<get<0>(TileShape_{}) / 128>;
static_assert(NumSplitsM{} == 1 || NumSplitsM{} == 2);
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128, _0> {
@@ -687,35 +689,37 @@ struct CollectiveMma<
template<
class AccumSlice,
class EngineAccum,
class LayoutAccum,
class ScaleFactor
>
CUTLASS_DEVICE
void scale_if_needed(GmmaFP8Accumulation<EngineAccum, LayoutAccum>& accumulation, ScaleFactor scaleFactor) {
void scale_if_needed(AccumSlice & accum, GmmaFP8Accumulation<EngineAccum, LayoutAccum>& accumulation, ScaleFactor scaleFactor) {
if constexpr (ScalePromotionInterval != 4) {
accumulation.scale_if_needed(scaleFactor);
accumulation.scale_if_needed(accum, scaleFactor);
}
else {
// avoid unnecessary tests when granularity is the finnest
accumulation.scale(scaleFactor);
accumulation.scale(accum, scaleFactor);
}
}
template<
class AccumSlice,
class EngineAccum,
class LayoutAccum,
class ScaleFactor1,
class ScaleFactor2
>
CUTLASS_DEVICE
void scale_if_needed(GmmaFP8Accumulation<EngineAccum, LayoutAccum>& accumulation, ScaleFactor1 scaleFactor1, ScaleFactor2 scaleFactor2) {
void scale_if_needed(AccumSlice & accum, GmmaFP8Accumulation<EngineAccum, LayoutAccum>& accumulation, ScaleFactor1 scaleFactor1, ScaleFactor2 scaleFactor2) {
if constexpr (ScalePromotionInterval != 4) {
accumulation.scale_if_needed(scaleFactor1, scaleFactor2);
accumulation.scale_if_needed(accum, scaleFactor1, scaleFactor2);
}
else {
// avoid unnecessary tests when granularity is the finnest
accumulation.scale(scaleFactor1, scaleFactor2);
accumulation.scale(accum, scaleFactor1, scaleFactor2);
}
}
@@ -821,75 +825,26 @@ struct CollectiveMma<
// Prologue GMMAs
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
// Tile accum
using NumSplitsM_Scale = cute::conditional_t<ScaleMsPerTile == 1, _1, NumSplitsM>;
static constexpr int ScaleMsPerWave = ScaleMsPerTile == 1 ? 1 : ScaleMsPerTile / NumSplitsM{};
auto accum_tiled = tiled_divide(accum, cute::tuple<_1, NumSplitsM>{});
auto tCrA_tiled = tiled_divide(tCrA, cute::tuple<_1, NumSplitsM>{});
auto tCsSFA_tiled = tiled_divide(tCsSFA, cute::tuple<_1, NumSplitsM_Scale>{});
auto tCrSFA_tiled = tiled_divide(tCrSFA, cute::tuple<_1, NumSplitsM_Scale>{});
auto tCrSFB_tiled = tiled_divide(tCrSFB, cute::tuple<_1, NumSplitsM_Scale>{});
// Temporary accumulator used by MMA
// On promotion, accumulated values are scaled and copied into `accum`
auto accum_temp = cute::make_fragment_like(accum_tiled(_0{}, _, _, _));
auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
GmmaFP8Accumulation accumulation(accum, ScalePromotionInterval, size<2>(tCrA));
warpgroup_fence_operand(accumulation());
if (k_tile_count > 0) {
// WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
pipeline.consumer_wait(smem_pipe_read, barrier_token);
int read_stage = smem_pipe_read.index();
// Load per block scale values from shared memory to registers
copy(tCsSFA(_,_,_,make_coord(_0{},read_stage)), tCrSFA);
copy(tCsSFB(_,_,_,make_coord(_0{},read_stage)), tCrSFB);
warpgroup_fence_operand(accumulation());
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
}
warpgroup_commit_batch();
warpgroup_fence_operand(accumulation());
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{});
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_b = tCrSFB(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) {
filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b;
}
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
ElementBlockScale scale_a = tCrSFA(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) {
filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a;
}
}
warpgroup_wait<0>();
++smem_pipe_read;
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
// Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB`
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA(_0{});
scale_if_needed(accumulation, scale_ab);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
scale_if_needed(accumulation, tCrSFA);
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFB);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFA, tCrSFB);
}
}
// Secondary accumulator for FP32 accum
GmmaFP8Accumulation accumulation(accum_temp, ScalePromotionInterval, size<2>(tCrA));
warpgroup_fence_operand(accumulation());
// Mainloop GMMAs
k_tile_count--;
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 1; --k_tile_count)
@@ -903,73 +858,86 @@ struct CollectiveMma<
int read_stage = smem_pipe_read.index();
// Load per block scale values from shared memory to registers (at most twice per block along M and/or N)
copy(tCsSFA(_,_,_,make_coord(_0{}, read_stage)), tCrSFA);
copy(tCsSFB(_,_,_,make_coord(_0{}, read_stage)), tCrSFB);
copy(tCsSFB(_,_,_,make_coord(_0{}, read_stage)), tCrSFB);
if constexpr (ScalePromotionInterval != 4) {
if (accumulation.prepare_if_needed()) {
CUTLASS_PRAGMA_UNROLL
for (int m_split = 0; m_split < NumSplitsM{}; ++m_split) {
auto tCrA_local = tCrA_tiled(m_split, _, _, _, _);
auto tCrSFA_local = tCrSFA_tiled(m_split, _, _, _);
auto tCrSFB_local = tCrSFB_tiled(m_split, _, _, _);
auto accum_local = accum_tiled(m_split, _, _, _);
copy(tCsSFA_tiled(m_split, _, _, _, make_coord(_0{}, read_stage)), tCrSFA_local);
bool is_last = (m_split == NumSplitsM{} - 1);
if constexpr (ScalePromotionInterval != 4) {
if (accumulation.prepare_if_needed()) {
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
}
else {
// Always zero out the accumulator for finest granularity
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
}
else {
// Always zero out the accumulator for finest granularity
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
warpgroup_fence_operand(accumulation());
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
}
warpgroup_commit_batch();
/// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed
warpgroup_fence_operand(accumulation());
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{});
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_b = tCrSFB(_0{});
warpgroup_fence_operand(accumulation());
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) {
filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b;
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_local(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
}
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
ElementBlockScale scale_a = tCrSFA(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) {
filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a;
warpgroup_commit_batch();
/// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed
warpgroup_fence_operand(accumulation());
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) {
tCrSFA_local(_0{}) = tCrSFA_local(_0{}) * tCrSFB(_0{});
}
if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_b = tCrSFB(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFA_local)); i++) {
filter_zeros(tCrSFA_local)(i) = filter_zeros(tCrSFA_local)(i) * scale_b;
}
}
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) {
ElementBlockScale scale_a = tCrSFA_local(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFB_local)); i++) {
filter_zeros(tCrSFB_local)(i) = filter_zeros(tCrSFB_local)(i) * scale_a;
}
}
}
warpgroup_wait<0>();
pipeline.consumer_release(smem_pipe_release); // Unlock previous tile
++smem_pipe_read;
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
warpgroup_wait<0>();
if (is_last) {
pipeline.consumer_release(smem_pipe_release); // Unlock previous tile
++smem_pipe_read;
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
}
// Block scale the accumulators with reg tensor `tCrSFA_local` and `tCrSFB`
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA_local(_0{});
scale_if_needed(accum_local, accumulation, scale_ab);
}
if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) {
scale_if_needed(accum_local, accumulation, tCrSFA_local);
}
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) {
scale_if_needed(accum_local, accumulation, tCrSFB_local);
}
if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile > 1) {
scale_if_needed(accum_local, accumulation, tCrSFA_local, tCrSFB_local);
}
// Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB`
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA(_0{});
scale_if_needed(accumulation, scale_ab);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
scale_if_needed(accumulation, tCrSFA);
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFB);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFA, tCrSFB);
}
// Advance smem_pipe_read and smem_pipe_release
++smem_pipe_release;
if (is_last) {
// Advance smem_pipe_read and smem_pipe_release
++smem_pipe_release;
}
} // end for (m_split)
}
if (k_tile_count > 0) {
@@ -981,97 +949,101 @@ struct CollectiveMma<
int read_stage = smem_pipe_read.index();
// Load per block scale values from shared memory to registers (at most twice per block along M and/or N)
copy(tCsSFA(_,_,_,make_coord(_0{}, read_stage)), tCrSFA);
copy(tCsSFB(_,_,_,make_coord(_0{}, read_stage)), tCrSFB);
CUTLASS_PRAGMA_UNROLL
for (int m_split = 0; m_split < NumSplitsM{}; ++m_split) {
auto tCrA_local = tCrA_tiled(m_split, _, _, _, _);
auto tCrSFA_local = tCrSFA_tiled(m_split, _, _, _);
auto tCrSFB_local = tCrSFB_tiled(m_split, _, _, _);
auto accum_local = accum_tiled(m_split, _, _, _);
copy(tCsSFA_tiled(m_split, _, _, _, make_coord(_0{}, read_stage)), tCrSFA_local);
bool is_last = (m_split == NumSplitsM{} - 1);
if constexpr (ScalePromotionInterval != 4) {
if (accumulation.prepare_if_needed()) {
if constexpr (ScalePromotionInterval != 4) {
if (accumulation.prepare_if_needed()) {
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
}
else {
// Always zero out the accumulator for finest granularity
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
}
else {
// Always zero out the accumulator for finest granularity
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
warpgroup_fence_operand(accumulation());
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
}
warpgroup_commit_batch();
/// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed
warpgroup_fence_operand(accumulation());
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{});
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_b = tCrSFB(_0{});
warpgroup_fence_operand(accumulation());
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) {
filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b;
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_local(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
}
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
ElementBlockScale scale_a = tCrSFA(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) {
filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a;
warpgroup_commit_batch();
/// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed
warpgroup_fence_operand(accumulation());
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) {
tCrSFA_local(_0{}) = tCrSFA_local(_0{}) * tCrSFB(_0{});
}
if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_b = tCrSFB(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFA_local)); i++) {
filter_zeros(tCrSFA_local)(i) = filter_zeros(tCrSFA_local)(i) * scale_b;
}
}
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) {
ElementBlockScale scale_a = tCrSFA_local(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFB_local)); i++) {
filter_zeros(tCrSFB_local)(i) = filter_zeros(tCrSFB_local)(i) * scale_a;
}
}
warpgroup_wait<0>();
if (is_last) {
pipeline.consumer_release(smem_pipe_release); // Unlock previous tile
}
// Block scale the accumulators with reg tensor `tCrSFA_local` and `tCrSFB`
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA_local(_0{});
scale_if_needed(accum_local, accumulation, scale_ab);
}
if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) {
scale_if_needed(accum_local, accumulation, tCrSFA_local);
}
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) {
scale_if_needed(accum_local, accumulation, tCrSFB_local);
}
if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile > 1) {
scale_if_needed(accum_local, accumulation, tCrSFA_local, tCrSFB_local);
}
if constexpr (ScalePromotionInterval != 4) {
// residues only exists when granularity is not the finnest
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA_local(_0{});
accumulation.scale_residue_if_needed(accum_local, scale_ab);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
accumulation.scale_residue_if_needed(accum_local, tCrSFA_local);
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
accumulation.scale_residue_if_needed(accum_local, tCrSFB_local);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
accumulation.scale_residue_if_needed(accum_local, tCrSFA_local, tCrSFB_local);
}
}
}
warpgroup_wait<0>();
pipeline.consumer_release(smem_pipe_release); // Unlock previous tile
// Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB`
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA(_0{});
scale_if_needed(accumulation, scale_ab);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
scale_if_needed(accumulation, tCrSFA);
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFB);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFA, tCrSFB);
}
warpgroup_fence_operand(accumulation());
} // end for (m_split)
}
if constexpr (ScalePromotionInterval != 4) {
// residues only exists when granularity is not the finnest
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA(_0{});
accumulation.scale_residue_if_needed(scale_ab);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
accumulation.scale_residue_if_needed(tCrSFA);
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
accumulation.scale_residue_if_needed(tCrSFB);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
accumulation.scale_residue_if_needed(tCrSFA, tCrSFB);
}
}
warpgroup_fence_operand(accumulation());
}
/// Perform a Consumer Epilogue to release all buffers
CUTLASS_DEVICE void
mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) {
if (k_tile_count > 0) {
// The pipeline is not released in the first iteration
smem_pipe_release.advance(k_tile_count - 1);
pipeline.consumer_release(smem_pipe_release);
}
}
//
include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp
+5 -5
View File
@@ -481,8 +481,8 @@ struct CollectiveMma<
int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
GmmaFP8Accumulation accumulation(accum, mainloop_params.mma_promotion_interval, size<2>(tCrA));
auto accm_temp = cute::make_fragment_like(accum);
GmmaFP8Accumulation accumulation(accm_temp, mainloop_params.mma_promotion_interval, size<2>(tCrA));
warpgroup_fence_operand(accumulation());
CUTLASS_PRAGMA_UNROLL
for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0; --k_tile_prologue)
@@ -506,7 +506,7 @@ struct CollectiveMma<
}
warpgroup_commit_batch();
accumulation.promote_if_needed();
accumulation.promote_if_needed(accum);
++smem_pipe_read;
}
@@ -547,7 +547,7 @@ struct CollectiveMma<
warpgroup_wait<K_PIPE_MMAS>();
warpgroup_fence_operand(accumulation());
accumulation.promote_if_needed();
accumulation.promote_if_needed(accum);
pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it
@@ -556,7 +556,7 @@ struct CollectiveMma<
++smem_pipe_release;
}
accumulation.promote_residue_if_needed();
accumulation.promote_residue_if_needed(accum);
warpgroup_fence_operand(accumulation());
}
include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+181 -194
View File
@@ -33,7 +33,9 @@
#include "cutlass/cutlass.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/fp8_accumulation.hpp"
#include "cutlass/trace.h"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/numeric_types.h"
#include "cute/arch/cluster_sm90.hpp"
@@ -203,6 +205,9 @@ struct CollectiveMma<
static_assert(cute::is_same_v<ElementAccumulator, ElementBlockScale>,
"ElementAccumulator and ElementBlockScale should be same datatype");
using NumSplitsM = cute::C<get<0>(TileShape_{}) / 128>;
static_assert(NumSplitsM{} == 1 || NumSplitsM{} == 2);
struct SharedStorage
{
struct TensorStorage : cute::aligned_struct<128> {
@@ -720,34 +725,36 @@ struct CollectiveMma<
}
template<
class AccumSlice,
class EngineAccum,
class LayoutAccum,
class ScaleFactor
>
CUTLASS_DEVICE
void scale_if_needed(GmmaFP8Accumulation<EngineAccum, LayoutAccum>& accumulation, ScaleFactor scaleFactor) {
void scale_if_needed(AccumSlice & accum, GmmaFP8Accumulation<EngineAccum, LayoutAccum>& accumulation, ScaleFactor scaleFactor) {
if constexpr (ScalePromotionInterval != 4) {
accumulation.scale_if_needed(scaleFactor);
accumulation.scale_if_needed(accum, scaleFactor);
}
else {
// avoid unnecessary tests when granularity is the finnest
accumulation.scale(scaleFactor);
accumulation.scale(accum, scaleFactor);
}
}
template<
class AccumSlice,
class EngineAccum,
class LayoutAccum,
class ScaleFactor1,
class ScaleFactor2
>
CUTLASS_DEVICE
void scale_if_needed(GmmaFP8Accumulation<EngineAccum, LayoutAccum>& accumulation, ScaleFactor1 scaleFactor1, ScaleFactor2 scaleFactor2) {
void scale_if_needed(AccumSlice & accum, GmmaFP8Accumulation<EngineAccum, LayoutAccum>& accumulation, ScaleFactor1 scaleFactor1, ScaleFactor2 scaleFactor2) {
if constexpr (ScalePromotionInterval != 4) {
accumulation.scale_if_needed(scaleFactor1, scaleFactor2);
accumulation.scale_if_needed(accum, scaleFactor1, scaleFactor2);
}
else {
// avoid unnecessary tests when granularity is the finnest
accumulation.scale(scaleFactor1, scaleFactor2);
accumulation.scale(accum, scaleFactor1, scaleFactor2);
}
}
@@ -852,70 +859,27 @@ struct CollectiveMma<
// Prologue GMMAs
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
// Tile accum
using NumSplitsM_Scale = cute::conditional_t<ScaleMsPerTile == 1, _1, NumSplitsM>;
static constexpr int ScaleMsPerWave = ScaleMsPerTile == 1 ? 1 : ScaleMsPerTile / NumSplitsM{};
auto accum_tiled = tiled_divide(accum, cute::tuple<_1, NumSplitsM>{});
auto tCrA_tiled = tiled_divide(tCrA, cute::tuple<_1, NumSplitsM>{});
auto tCsSFA_tiled = tiled_divide(tCsSFA, cute::tuple<_1, NumSplitsM_Scale>{});
auto tCrSFA_tiled = tiled_divide(tCrSFA, cute::tuple<_1, NumSplitsM_Scale>{});
auto tCrSFB_tiled = tiled_divide(tCrSFB, cute::tuple<_1, NumSplitsM_Scale>{});
// Temporary accumulator used by MMA
// On promotion, accumulated values are scaled and copied into `accum`
auto accum_temp = cute::make_fragment_like(accum_tiled(_0{}, _, _, _));
// WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
pipeline.consumer_wait(smem_pipe_read, barrier_token);
GmmaFP8Accumulation accumulation(accum, ScalePromotionInterval, size<2>(tCrA));
warpgroup_fence_operand(accumulation());
{
int read_stage = smem_pipe_read.index();
// Load per block scale values from shared memory to registers
copy(tCsSFA(_,_,_,make_coord(_0{}, read_stage)), tCrSFA);
copy(tCsSFB(_,_,_,make_coord(_0{}, read_stage)), tCrSFB);
warpgroup_fence_operand(accumulation());
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
}
warpgroup_commit_batch();
warpgroup_fence_operand(accumulation());
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{});
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_b = tCrSFB(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) {
filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b;
}
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
ElementBlockScale scale_a = tCrSFA(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) {
filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a;
}
}
warpgroup_wait<0>();
++smem_pipe_read;
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
// Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB`
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA(_0{});
scale_if_needed(accumulation, scale_ab);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
scale_if_needed(accumulation, tCrSFA);
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFB);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFA, tCrSFB);
}
}
// Secondary accumulator for FP32 accum
GmmaFP8Accumulation accumulation(accum_temp, ScalePromotionInterval, size<2>(tCrA));
warpgroup_fence_operand(accumulation());
// Mainloop GMMAs
k_tile_count -= 1;
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 1; --k_tile_count)
@@ -929,71 +893,86 @@ struct CollectiveMma<
int read_stage = smem_pipe_read.index();
// Load per block scale values from shared memory to registers (at most twice per block along M and/or N)
copy(tCsSFA(_,_,_,make_coord(_0{}, read_stage)), tCrSFA);
copy(tCsSFB(_,_,_,make_coord(_0{}, read_stage)), tCrSFB);
if constexpr (ScalePromotionInterval != 4) {
if (accumulation.prepare_if_needed()) {
CUTLASS_PRAGMA_UNROLL
for (int m_split = 0; m_split < NumSplitsM{}; ++m_split) {
auto tCrA_local = tCrA_tiled(m_split, _, _, _, _);
auto tCrSFA_local = tCrSFA_tiled(m_split, _, _, _);
auto tCrSFB_local = tCrSFB_tiled(m_split, _, _, _);
auto accum_local = accum_tiled(m_split, _, _, _);
copy(tCsSFA_tiled(m_split, _, _, _, make_coord(_0{}, read_stage)), tCrSFA_local);
bool is_last = (m_split == NumSplitsM{} - 1);
if constexpr (ScalePromotionInterval != 4) {
if (accumulation.prepare_if_needed()) {
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
}
else {
// Always zero out the accumulator for finest granularity
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
}
else {
// Always zero out the accumulator for finest granularity
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
warpgroup_fence_operand(accumulation());
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
}
warpgroup_commit_batch();
/// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed
warpgroup_fence_operand(accumulation());
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{});
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_b = tCrSFB(_0{});
warpgroup_fence_operand(accumulation());
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) {
filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b;
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_local(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
}
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
ElementBlockScale scale_a = tCrSFA(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) {
filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a;
}
}
warpgroup_wait<0>();
pipeline.consumer_release(smem_pipe_release); // Unlock previous tile
++smem_pipe_read;
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
// Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB`
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA(_0{});
scale_if_needed(accumulation, scale_ab);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
scale_if_needed(accumulation, tCrSFA);
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFB);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFA, tCrSFB);
}
warpgroup_commit_batch();
// Advance smem_pipe_read and smem_pipe_release
++smem_pipe_release;
/// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed
warpgroup_fence_operand(accumulation());
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) {
tCrSFA_local(_0{}) = tCrSFA_local(_0{}) * tCrSFB(_0{});
}
if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_b = tCrSFB(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFA_local)); i++) {
filter_zeros(tCrSFA_local)(i) = filter_zeros(tCrSFA_local)(i) * scale_b;
}
}
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) {
ElementBlockScale scale_a = tCrSFA_local(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFB_local)); i++) {
filter_zeros(tCrSFB_local)(i) = filter_zeros(tCrSFB_local)(i) * scale_a;
}
}
warpgroup_wait<0>();
if (is_last) {
pipeline.consumer_release(smem_pipe_release); // Unlock previous tile
++smem_pipe_read;
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
}
// Block scale the accumulators with reg tensor `tCrSFA_local` and `tCrSFB`
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA_local(_0{});
scale_if_needed(accum_local, accumulation, scale_ab);
}
if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) {
scale_if_needed(accum_local, accumulation, tCrSFA_local);
}
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) {
scale_if_needed(accum_local, accumulation, tCrSFB_local);
}
if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile > 1) {
scale_if_needed(accum_local, accumulation, tCrSFA_local, tCrSFB_local);
}
if (is_last) {
// Advance smem_pipe_read and smem_pipe_release
++smem_pipe_release;
}
} // end for (m_split)
}
if (k_tile_count) {
pipeline.consumer_wait(smem_pipe_read, barrier_token);
@@ -1005,93 +984,101 @@ struct CollectiveMma<
int read_stage = smem_pipe_read.index();
// Load per block scale values from shared memory to registers (at most twice per block along M and/or N)
copy(tCsSFA(_,_,_,make_coord(_0{}, read_stage)), tCrSFA);
copy(tCsSFB(_,_,_,make_coord(_0{}, read_stage)), tCrSFB);
CUTLASS_PRAGMA_UNROLL
for (int m_split = 0; m_split < NumSplitsM{}; ++m_split) {
auto tCrA_local = tCrA_tiled(m_split, _, _, _, _);
auto tCrSFA_local = tCrSFA_tiled(m_split, _, _, _);
auto tCrSFB_local = tCrSFB_tiled(m_split, _, _, _);
auto accum_local = accum_tiled(m_split, _, _, _);
copy(tCsSFA_tiled(m_split, _, _, _, make_coord(_0{}, read_stage)), tCrSFA_local);
bool is_last = (m_split == NumSplitsM{} - 1);
if constexpr (ScalePromotionInterval != 4) {
if (accumulation.prepare_if_needed()) {
if constexpr (ScalePromotionInterval != 4) {
if (accumulation.prepare_if_needed()) {
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
}
else {
// Always zero out the accumulator for finest granularity
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
}
else {
// Always zero out the accumulator for finest granularity
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
}
warpgroup_fence_operand(accumulation());
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
}
warpgroup_commit_batch();
/// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed
warpgroup_fence_operand(accumulation());
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
tCrSFA(_0{}) = tCrSFA(_0{}) * tCrSFB(_0{});
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_b = tCrSFB(_0{});
warpgroup_fence_operand(accumulation());
warpgroup_arrive();
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFA)); i++) {
filter_zeros(tCrSFA)(i) = filter_zeros(tCrSFA)(i) * scale_b;
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_local(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
}
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
ElementBlockScale scale_a = tCrSFA(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFB)); i++) {
filter_zeros(tCrSFB)(i) = filter_zeros(tCrSFB)(i) * scale_a;
}
}
warpgroup_wait<0>();
pipeline.consumer_release(smem_pipe_release); // Unlock previous tile
// Block scale the accumulators with reg tensor `tCrSFA` and `tCrSFB`
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA(_0{});
scale_if_needed(accumulation, scale_ab);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
scale_if_needed(accumulation, tCrSFA);
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFB);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
scale_if_needed(accumulation, tCrSFA, tCrSFB);
}
}
if constexpr (ScalePromotionInterval != 4) {
// residues only exists when granularity is not the finnest
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA(_0{});
accumulation.scale_residue_if_needed(scale_ab);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
accumulation.scale_residue_if_needed(tCrSFA);
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
accumulation.scale_residue_if_needed(tCrSFB);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
accumulation.scale_residue_if_needed(tCrSFA, tCrSFB);
}
}
warpgroup_commit_batch();
warpgroup_fence_operand(accumulation());
/// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed
warpgroup_fence_operand(accumulation());
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) {
tCrSFA_local(_0{}) = tCrSFA_local(_0{}) * tCrSFB(_0{});
}
if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_b = tCrSFB(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFA_local)); i++) {
filter_zeros(tCrSFA_local)(i) = filter_zeros(tCrSFA_local)(i) * scale_b;
}
}
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) {
ElementBlockScale scale_a = tCrSFA_local(_0{});
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(filter_zeros(tCrSFB_local)); i++) {
filter_zeros(tCrSFB_local)(i) = filter_zeros(tCrSFB_local)(i) * scale_a;
}
}
warpgroup_wait<0>();
if (is_last) {
pipeline.consumer_release(smem_pipe_release); // Unlock previous tile
}
// Block scale the accumulators with reg tensor `tCrSFA_local` and `tCrSFB`
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA_local(_0{});
scale_if_needed(accum_local, accumulation, scale_ab);
}
if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile == 1) {
scale_if_needed(accum_local, accumulation, tCrSFA_local);
}
if constexpr (ScaleMsPerWave == 1 && ScaleNsPerTile > 1) {
scale_if_needed(accum_local, accumulation, tCrSFB_local);
}
if constexpr (ScaleMsPerWave > 1 && ScaleNsPerTile > 1) {
scale_if_needed(accum_local, accumulation, tCrSFA_local, tCrSFB_local);
}
if constexpr (ScalePromotionInterval != 4) {
// residues only exists when granularity is not the finnest
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
ElementBlockScale scale_ab = tCrSFA_local(_0{});
accumulation.scale_residue_if_needed(accum_local, scale_ab);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
accumulation.scale_residue_if_needed(accum_local, tCrSFA_local);
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
accumulation.scale_residue_if_needed(accum_local, tCrSFB_local);
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
accumulation.scale_residue_if_needed(accum_local, tCrSFA_local, tCrSFB_local);
}
}
warpgroup_fence_operand(accumulation());
} // end for (m_split)
}
}
/// Perform a Consumer Epilogue to release all buffers
CUTLASS_DEVICE void
mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) {
// The pipeline is not released in the first iteration
smem_pipe_release.advance(k_tile_count - 1);
pipeline.consumer_release(smem_pipe_release);
}
};
include/cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized_fp8.hpp
+5 -5
View File
@@ -586,8 +586,8 @@ struct CollectiveMma<
int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
GmmaFP8Accumulation accumulation(accum, mainloop_params.mma_promotion_interval, size<2>(tCrA));
auto accm_temp = cute::make_fragment_like(accum);
GmmaFP8Accumulation accumulation(accm_temp, mainloop_params.mma_promotion_interval, size<2>(tCrA));
warpgroup_fence_operand(accumulation());
CUTLASS_PRAGMA_UNROLL
for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0; --k_tile_prologue)
@@ -614,7 +614,7 @@ struct CollectiveMma<
warpgroup_commit_batch();
accumulation.promote_if_needed();
accumulation.promote_if_needed(accum);
++smem_pipe_read;
}
@@ -652,7 +652,7 @@ struct CollectiveMma<
warpgroup_wait<K_PIPE_MMAS>();
warpgroup_fence_operand(accumulation());
accumulation.promote_if_needed();
accumulation.promote_if_needed(accum);
// UNLOCK smem_pipe_release, done _computing_ on it
pipeline.consumer_release(smem_pipe_release);
@@ -662,7 +662,7 @@ struct CollectiveMma<
++smem_pipe_release;
}
accumulation.promote_residue_if_needed();
accumulation.promote_residue_if_needed(accum);
warpgroup_fence_operand(accumulation());
}
include/cutlass/gemm/group_array_problem_shape.hpp
+50 -7
View File
@@ -44,6 +44,7 @@
#include <initializer_list>
#endif
////////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm {
@@ -57,6 +58,7 @@ struct GroupProblemShape {
UnderlyingProblemShape* problem_shapes = nullptr;
UnderlyingProblemShape const* host_problem_shapes = nullptr;
CUTLASS_HOST_DEVICE
int32_t groups() const { return num_groups; }
@@ -79,16 +81,56 @@ struct GroupProblemShape {
}
};
template <class ProblemShape_, class MaxProblemShape_>
template <class ProblemShape_>
struct MoEProblemShape {
using UnderlyingProblemShape = ProblemShape_;
using MaxProblemShape = MaxProblemShape_;
static_assert(rank(UnderlyingProblemShape{}) == 3, "ProblemShape{} should be <M,N,K>");
int32_t max_m = 0;
int32_t max_n = 0;
int32_t max_k = 0;
int32_t num_groups = 0;
int32_t* tokens_per_expert = nullptr;
int32_t* tokens_per_expert_host = nullptr;
CUTLASS_HOST_DEVICE
int32_t groups() const { return num_groups; }
CUTLASS_HOST_DEVICE
UnderlyingProblemShape const
get_problem_shape(int32_t group_idx=0) const {
UnderlyingProblemShape expert_problem_dims;
assert(tokens_per_expert != nullptr); //tokens_per_expert should not be null
if (group_idx < num_groups) { // add check on the can_implement
expert_problem_dims = {max_m, tokens_per_expert[group_idx], max_k};
}
return expert_problem_dims;
}
// Function returns max problem shape if tokens_per_expert host is unavailable.
// Returns host problem shape if tokens_per_expert host is available.
CUTLASS_HOST_DEVICE
UnderlyingProblemShape const
get_host_problem_shape(int32_t group_idx=0) const {
UnderlyingProblemShape expert_problem_dims = {max_m, max_n, max_k};
assert(tokens_per_expert_host != nullptr); //tokens_per_expert_host should not be null
if (group_idx < num_groups) {
expert_problem_dims = {max_m, tokens_per_expert_host[group_idx], max_k};
}
return expert_problem_dims;
}
CUTLASS_HOST_DEVICE
bool
is_host_problem_shape_available() const {
return tokens_per_expert_host != nullptr;
}
UnderlyingProblemShape problem_shape;
MaxProblemShape max_problem_shape;
};
template <class ProblemShape_>
class ArrayProblemShape {
public:
@@ -135,8 +177,9 @@ namespace detail {
template<class T>
struct is_moe_problem_shape : cute::false_type {};
template<class T, class U>
struct is_moe_problem_shape<cutlass::gemm::MoEProblemShape<T,U>> : cute::true_type {};
template<class T>
struct is_moe_problem_shape<cutlass::gemm::MoEProblemShape<T>> : cute::true_type {};
}
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp
+6 -7
View File
@@ -842,13 +842,12 @@ public:
);
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();
}
// 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;
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp
+7 -4
View File
@@ -649,18 +649,21 @@ public:
transform2mma_pipeline.init_masks(cluster_shape);
mma2accum_pipeline.init_masks(cluster_shape);
// Allocate accumulators
auto acc_shape = collective_mainloop.partition_accumulator_shape();
// 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);
// Ensure memory ops in this kernel are not done prior to completion of dependent grids.
cutlass::arch::wait_on_dependent_grids();
typename TileScheduler::WorkTileInfo work_tile_info = scheduler.initial_work_tile_info(cluster_shape);
auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
// 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);
// Allocate accumulators
auto acc_shape = collective_mainloop.partition_accumulator_shape();
// NOTE: we can assume the tmem buf starts at zero since we allocate all tmem in this kernel
auto bulk_tmem = TiledMma::make_fragment_C(append(acc_shape,
Int<AccumulatorPipelineStageCount>{}));
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_mma_transform.hpp
+9 -5
View File
@@ -726,11 +726,6 @@ public:
mainloop_ab_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);
auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
//
// TMEM "Allocation"
//
@@ -740,6 +735,15 @@ public:
Tensor accumulators = cutlass::detail::make_sm100_accumulator<AccumulatorPipelineStageCount, IsOverlappingAccum>(
tiled_mma, acc_shape, EpilogueTile{});
// TileID scheduler
TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster);
// Ensure memory ops in this kernel are not done prior to completion of dependent grids.
cutlass::arch::wait_on_dependent_grids();
typename TileScheduler::WorkTileInfo work_tile_info = scheduler.initial_work_tile_info(cluster_shape);
auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
pipeline_init_wait(cluster_size);
if constexpr (IsGroupedGemmKernel) {
include/cutlass/gemm/kernel/sm100_gemm_mixed_tma_cpasync_warpspecialized.hpp
+33 -44
View File
@@ -83,20 +83,20 @@ public:
CUTLASS_HOST_DEVICE
static auto get_problem_shape_gemm(ProblemShape const& shape) {
if constexpr (IsGroupedGemmKernel) {
return shape.max_problem_shape;
auto problem_shape_MNK = shape.get_host_problem_shape(0); //gets the maximum problem shape here
auto problem_shape_MNKL = append<4>(problem_shape_MNK, shape.groups()); //appends num_groups to it
return problem_shape_MNKL;
}
else {
return shape;
}
}
CUTLASS_HOST_DEVICE
static auto get_problem_shape_scheduler(ProblemShape const& shape) {
if constexpr (IsMoEScheduler) {
if constexpr (IsMoEScheduler) {
return shape;
}
else if constexpr (IsGroupedGemmKernel) {
return shape.problem_shape;
}
else {
return shape;
}
@@ -106,7 +106,7 @@ public:
CUTLASS_HOST_DEVICE
static auto get_effective_shape(ProblemShape const& shape, WorkTileInfo const& work_tile_info) {
if constexpr (IsGroupedGemmKernel) {
return append<4>(shape.problem_shape.get_problem_shape(work_tile_info.L_idx), Int<1>{});
return append<4>(shape.get_problem_shape(work_tile_info.L_idx), Int<1>{});
}
else {
return append<4>(shape, Int<1>{});
@@ -114,10 +114,9 @@ public:
}
using ProblemShapeGemm = decltype(get_problem_shape_gemm(ProblemShape{}));
using ProblemShapeScheduler = decltype(get_problem_shape_scheduler(ProblemShape{}));
static_assert(rank(ProblemShapeGemm{}) == 3 or rank(ProblemShapeGemm{}) == 4,
"ProblemShapeGemm{} should be <M,N,K> or <M,N,K,L>");
"ProblemShape{} should be <M,N,K> or <M,N,K,L>");
static constexpr bool IsGdcEnabled = false;
// Mainloop derived types
using CollectiveMainloop = CollectiveMainloop_;
@@ -160,7 +159,7 @@ public:
static_assert(size(AtomThrShapeMNK{}) == 1, "Lower alignment kernel only supports 1x1x1 cluster shape.");
using TileSchedulerTag = cute::conditional_t<IsGroupedGemmKernel && !IsMoEScheduler, GroupScheduler, TileSchedulerTag_>;
using TileScheduler = typename detail::TileSchedulerSelector<
TileSchedulerTag, ArchTag, CtaShape_MNK, ClusterShape, SchedulerPipelineStageCount, ProblemShapeScheduler>::Scheduler;
TileSchedulerTag, ArchTag, CtaShape_MNK, ClusterShape, SchedulerPipelineStageCount, ProblemShape>::Scheduler;
using TileSchedulerArguments = typename TileScheduler::Arguments;
using TileSchedulerParams = typename TileScheduler::Params;
@@ -259,7 +258,6 @@ public:
GemmUniversalMode mode{};
ProblemShape problem_shape{};
ProblemShapeGemm problem_shape_gemm{};
ProblemShapeScheduler problem_shape_scheduler{};
MainloopParams mainloop{};
EpilogueParams epilogue{};
KernelHardwareInfo hw_info{};
@@ -289,18 +287,20 @@ public:
Params
to_underlying_arguments(Arguments const& args, void* workspace) {
(void) workspace;
// auto problem_shape = args.problem_shape;
// auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto problem_shape = args.problem_shape;
auto problem_shape_gemm = get_problem_shape_gemm(args.problem_shape);
auto problem_shape_scheduler = get_problem_shape_scheduler(args.problem_shape);
// Get SM count if needed, otherwise use user supplied SM count
int sm_count = args.hw_info.sm_count;
if (sm_count != 0) {
if (IsGroupedGemmKernel && sm_count <= 0) {
CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n"
" For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
}
else if (!IsGroupedGemmKernel && sm_count != 0) {
CUTLASS_TRACE_HOST(" WARNING: SM100 tile scheduler does not allow for user specified SM counts.\n"
" To restrict a kernel's resource usage, consider using CUDA driver APIs instead (green contexts).");
sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
}
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
@@ -313,25 +313,24 @@ public:
// Epilogue
void* epilogue_workspace = workspace_ptr + workspace_offset;
workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
workspace_offset += CollectiveEpilogue::get_workspace_size(problem_shape, args.epilogue);
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
void* mainloop_workspace = nullptr;
// Tile scheduler
void* scheduler_workspace = workspace_ptr + workspace_offset;
workspace_offset += TileScheduler::template get_workspace_size<ProblemShapeScheduler, ElementAccumulator>(
args.scheduler, problem_shape_scheduler, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs);
workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
args.scheduler, problem_shape, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs);
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
TileSchedulerParams scheduler;
if constexpr (IsGroupedGemmKernel) {
scheduler = TileScheduler::to_underlying_arguments(
problem_shape_scheduler, TileShape{}, AtomThrShapeMNK{}, ClusterShape{},
problem_shape, TileShape{}, AtomThrShapeMNK{}, ClusterShape{},
args.hw_info, args.scheduler, scheduler_workspace);
}
else {
auto problem_shape = args.problem_shape;
auto problem_shape_MNKL = append<4>(problem_shape, 1);
scheduler = TileScheduler::to_underlying_arguments(
@@ -344,7 +343,6 @@ public:
args.mode,
args.problem_shape,
problem_shape_gemm,
problem_shape_scheduler,
CollectiveMainloop::to_underlying_arguments(problem_shape_gemm, args.mainloop, mainloop_workspace),
CollectiveEpilogue::to_underlying_arguments(problem_shape_gemm, args.epilogue, epilogue_workspace),
hw_info,
@@ -358,8 +356,7 @@ public:
if constexpr (IsGroupedGemmKernel) {
implementable &= args.mode == GemmUniversalMode::kGrouped;
implementable &= rank(ProblemShapeGemm{}) == 4;
implementable &= rank(typename ProblemShape::UnderlyingProblemShape::UnderlyingProblemShape{}) == 3;
implementable &= rank(typename ProblemShape::UnderlyingProblemShape{}) == 3;
}
else {
implementable &= (args.mode == GemmUniversalMode::kGemm) or
@@ -387,15 +384,14 @@ public:
size_t workspace_size = 0;
auto problem_shape_gemm = get_problem_shape_gemm(args.problem_shape);
auto problem_shape_scheduler = get_problem_shape_scheduler(args.problem_shape);
// Epilogue
workspace_size += CollectiveEpilogue::get_workspace_size(problem_shape_gemm, args.epilogue);
workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
// Tile scheduler
workspace_size += TileScheduler::template get_workspace_size<ProblemShapeScheduler, ElementAccumulator>(
args.scheduler, problem_shape_scheduler, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs);
workspace_size += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
args.scheduler, args.problem_shape, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs);
workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
return workspace_size;
@@ -409,7 +405,6 @@ public:
size_t workspace_offset = 0;
auto problem_shape_gemm = get_problem_shape_gemm(args.problem_shape);
auto problem_shape_scheduler = get_problem_shape_scheduler(args.problem_shape);
// Epilogue
status = CollectiveEpilogue::initialize_workspace(problem_shape_gemm, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter);
@@ -421,10 +416,10 @@ public:
}
// Tile scheduler
status = TileScheduler::template initialize_workspace<ProblemShapeScheduler, ElementAccumulator>(
args.scheduler, workspace_ptr + workspace_offset, stream, problem_shape_scheduler, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs, cuda_adapter);
workspace_offset += TileScheduler::template get_workspace_size<ProblemShapeScheduler, ElementAccumulator>(
args.scheduler, problem_shape_scheduler, args.hw_info, NumFixupBarriers);
status = TileScheduler::template initialize_workspace<ProblemShape, ElementAccumulator>(
args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs, cuda_adapter);
workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
args.scheduler, args.problem_shape, args.hw_info, NumFixupBarriers);
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
if (status != Status::kSuccess) {
return status;
@@ -441,14 +436,14 @@ public:
if constexpr (IsGroupedGemmKernel) {
grid_shape = TileScheduler::get_grid_shape(
params.scheduler,
params.problem_shape_scheduler,
params.problem_shape,
TileShape{},
AtomThrShapeMNK{},
cluster_shape,
params.hw_info);
}
else {
auto problem_shape_MNKL = append<4>(params.problem_shape_scheduler, 1);
auto problem_shape_MNKL = append<4>(params.problem_shape, 1);
grid_shape = TileScheduler::get_grid_shape(
params.scheduler,
problem_shape_MNKL,
@@ -505,8 +500,6 @@ public:
// Do we load source tensor C or other aux inputs
bool is_epi_load_needed = collective_epilogue.is_producer_load_needed();
// printf("is_epi_load_needed = %d", (int)is_epi_load_needed);
IsParticipant is_participant = {
(warp_category == WarpCategory::MMA) && is_mma_leader_cta, // mma
(warp_category == WarpCategory::Sched) && is_first_cta_in_cluster, // sched
@@ -654,9 +647,6 @@ public:
//
// TMEM "Allocation"
//
// auto acc_shape = collective_mainloop.partition_accumulator_shape();
// auto bulk_tmem = TiledMma::make_fragment_C(append(acc_shape,
// Int<AccumulatorPipelineStageCount>{}));
auto tmem_storage = collective_mainloop.template init_tmem_tensors<EpilogueTile, IsOverlappingAccum>(EpilogueTile{});
//
@@ -666,10 +656,10 @@ public:
// Synchronization call. Blocks until barriers are initialized in shared memory.
pipeline_init_wait(cluster_size);
// __syncwarp();
// if (threadIdx.x % 32 == 0) {
// printf("warp %d start\n", warp_idx);
// }
if (not work_tile_info.is_valid()) {
// When problem shapes are only on device, the grid launched may be larger than the total number of blocks across groups
return;
}
if (is_participant.main_load_tma) {
// Ensure that the prefetched kernel does not touch
@@ -689,7 +679,6 @@ public:
// Get the number of K tiles to compute for this work as well as the starting K tile offset of the work.
auto k_tile_iter = scheduler.get_k_tile_iterator(work_tile_info, effective_shape, CtaShape_MNK{}, k_tiles);
auto k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, effective_shape, CtaShape_MNK{});
// auto k_tile_prologue = min(MainloopPipeline::Stages, k_tile_count);
auto [mainloop_producer_state_next_, unused_] = collective_mainloop.load_tma(
include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp
+9 -5
View File
@@ -734,11 +734,6 @@ public:
mainloop_ab_pipeline.init_masks(cluster_shape);
mainloop_sf_pipeline.init_masks(cluster_shape);
accumulator_pipeline.init_masks(cluster_shape);
// 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);
auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
//
// TMEM "Allocation"
//
@@ -749,6 +744,15 @@ public:
Tensor accumulators = cutlass::detail::make_sm100_accumulator<AccumulatorPipelineStageCount, IsOverlappingAccum>(
tiled_mma, acc_shape, EpilogueTile{});
// TileID scheduler
TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster);
// Ensure memory ops in this kernel are not done prior to completion of dependent grids.
cutlass::arch::wait_on_dependent_grids();
typename TileScheduler::WorkTileInfo work_tile_info = scheduler.initial_work_tile_info(cluster_shape);
auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
pipeline_init_wait(cluster_size);
if constexpr (IsGroupedGemmKernel) {
include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp
+5 -1
View File
@@ -48,6 +48,7 @@
#include "cutlass/trace.h"
#include "cutlass/gemm/kernel/sm90_tile_scheduler.hpp"
#include "cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp"
#include "cutlass/arch/grid_dependency_control.h"
///////////////////////////////////////////////////////////////////////////////
@@ -418,7 +419,7 @@ public:
// Any Tensor Op MMA Atom in the ISA is arch conditional.
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
#else
// Preconditions
@@ -581,6 +582,9 @@ public:
// Wait for all thread blocks in the Cluster
cluster_wait_fn();
// Ensure memory ops in this kernel are not done prior to completion of dependent grids.
cutlass::arch::wait_on_dependent_grids();
auto work_tile_info = scheduler.initial_work_tile_info(ClusterShape{});
if (not work_tile_info.is_valid()) {
include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp
+5 -1
View File
@@ -48,6 +48,7 @@
#include "cutlass/trace.h"
#include "cutlass/gemm/kernel/sm90_tile_scheduler.hpp"
#include "cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp"
#include "cutlass/arch/grid_dependency_control.h"
///////////////////////////////////////////////////////////////////////////////
@@ -430,7 +431,7 @@ public:
// Any Tensor Op MMA Atom in the ISA is arch conditional.
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
#else
// Preconditions
@@ -596,6 +597,9 @@ public:
// Wait for all thread blocks in the Cluster
cluster_wait_fn();
// Ensure memory ops in this kernel are not done prior to completion of dependent grids.
cutlass::arch::wait_on_dependent_grids();
auto work_tile_info = scheduler.initial_work_tile_info(ClusterShape{});
if (not work_tile_info.is_valid()) {
include/cutlass/gemm/kernel/sm90_gemm_tma.hpp
+6 -1
View File
@@ -42,6 +42,8 @@
#include "cutlass/gemm/kernel/sm90_tile_scheduler.hpp"
#include "cutlass/gemm/kernel/tile_scheduler.hpp"
#include "cutlass/trace.h"
#include "cutlass/arch/grid_dependency_control.h"
#include "cute/tensor.hpp"
///////////////////////////////////////////////////////////////////////////////
@@ -204,7 +206,7 @@ public:
// Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
#if ! defined(__CUDA_ARCH_FEAT_SM90_ALL)
printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
#else
// Preconditions
@@ -261,6 +263,9 @@ public:
auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA));
auto k_tile_count = size<2>(gA);
// Ensure memory ops in this kernel are not done prior to completion of dependent grids.
cutlass::arch::wait_on_dependent_grids();
// Perform the collective scoped MMA
CollectiveMainloop collective_mma;
collective_mma(
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp
+1 -1
View File
@@ -277,7 +277,7 @@ public:
// Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
#else
enum class WarpGroupRole {
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp
+16 -2
View File
@@ -132,7 +132,21 @@ public:
static constexpr int RegsPerThread =
size<0>(TileShape{}) * size<1>(TileShape{}) / NumMMAThreads *
sizeof(ElementAccumulator) / sizeof(uint32_t);
static constexpr bool HeavyRegisterPressure = RegsPerThread >= 208;
// Detect if this is SM120 blockscaled kernel which hits high register pressure
// on smaller tiles (e.g. 256x128 registers per thread)
template <typename T>
struct is_blockscaled : cute::false_type {};
template <int Stages, int SchedStages, class ClusterShape, class KernelSchedule>
struct is_blockscaled<MainloopSm120TmaWarpSpecializedBlockScaled<Stages, SchedStages, ClusterShape, KernelSchedule>>
: cute::true_type {};
static constexpr bool IsBlockScaled = is_blockscaled<DispatchPolicy>::value;
static constexpr bool HeavyRegisterPressure =
IsBlockScaled ? (RegsPerThread >= 128) : (RegsPerThread >= 208);
static constexpr uint32_t LoadRegisterRequirement = !HeavyRegisterPressure ? 40 : 24;
static constexpr uint32_t MmaRegisterRequirement = !HeavyRegisterPressure ? 232 : 240;
@@ -349,7 +363,7 @@ public:
// Any Tensor Op MMA Atom in the ISA is arch conditional.
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
#else
// Preconditions
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp
+1 -1
View File
@@ -361,7 +361,7 @@ public:
// Any Tensor Op MMA Atom in the ISA is arch conditional.
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
#else
// Preconditions
include/cutlass/gemm/kernel/sm90_gemm_warpspecialized.hpp
+6 -1
View File
@@ -42,6 +42,8 @@
#include "cutlass/gemm/kernel/sm90_tile_scheduler.hpp"
#include "cutlass/pipeline/pipeline.hpp"
#include "cute/tensor.hpp"
#include "cutlass/arch/grid_dependency_control.h"
///////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::kernel {
@@ -227,7 +229,7 @@ public:
// Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
#if ! defined(__CUDA_ARCH_FEAT_SM90_ALL)
printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
#else
enum class WarpGroupRole {
@@ -343,6 +345,9 @@ public:
auto k_residue = K - size<1>(gA) * size<2>(gA); // K - BLK_K * k_coord_max
auto residue_mnk = make_tuple(m_max_coord, n_max_coord, k_residue);
// Ensure memory ops in this kernel are not done prior to completion of dependent grids.
cutlass::arch::wait_on_dependent_grids();
collective_mainloop.load(
mainloop_pipeline,
mainloop_pipe_producer_state,
include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_cooperative.hpp
+5 -1
View File
@@ -42,6 +42,7 @@
#include "cutlass/gemm/kernel/tile_scheduler.hpp"
#include "cutlass/pipeline/pipeline.hpp"
#include "cute/tensor.hpp"
#include "cutlass/arch/grid_dependency_control.h"
///////////////////////////////////////////////////////////////////////////////
@@ -265,7 +266,7 @@ public:
// Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
#if ! defined(__CUDA_ARCH_FEAT_SM90_ALL)
printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
#else
static_assert(cute::rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>.");
@@ -386,6 +387,9 @@ public:
auto k_residue = K - size<1>(gA) * size<2>(gA); // K - BLK_K * k_coord_max
auto residue_mnk = make_tuple(m_max_coord, n_max_coord, k_residue);
// Ensure memory ops in this kernel are not done prior to completion of dependent grids.
cutlass::arch::wait_on_dependent_grids();
collective_mainloop.load(
mainloop_pipeline,
mainloop_pipe_producer_state,
include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_pingpong.hpp
+7 -1
View File
@@ -45,6 +45,8 @@
#include "cutlass/trace.h"
#include "cute/tensor.hpp"
#include "cutlass/arch/grid_dependency_control.h"
///////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::kernel {
@@ -271,7 +273,7 @@ public:
// Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
#if ! defined(__CUDA_ARCH_FEAT_SM90_ALL)
printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
#else
// Preconditions
@@ -409,6 +411,10 @@ public:
auto k_residue = K - size<1>(gA) * size<2>(gA); // K - BLK_K * k_coord_max
auto residue_mnk = make_tuple(m_max_coord, n_max_coord, k_residue);
// Ensure memory ops in this kernel are not done prior to completion of dependent grids.
cutlass::arch::wait_on_dependent_grids();
collective_mainloop.load(
mainloop_pipeline,
mainloop_pipe_producer_state,