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

support for K=0 for sm100 GG (#2746)

Browse Source
This commit is contained in:
ANIKET SHIVAM
2025-11-04 11:25:39 -05:00
committed by GitHub
parent 8afb19d904
commit 020c700e97
8 changed files with 72 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/epilogue/collective/detail.hpp
-3
View File
@@ -857,9 +857,6 @@ public:
TensorMap tensormap
)
{
// Wait for mma warp to fill tmem buffer with accumulator results
acc_pipeline.consumer_wait(acc_pipe_consumer_state);
auto [acc_state_next] = (*this).template operator()<ReuseTmem>(
acc_pipeline,
acc_pipe_consumer_state,
include/cutlass/epilogue/collective/sm100_epilogue_array_nosmem.hpp
+34 -17
View File
@@ -194,11 +194,15 @@ public:
static_assert(rank(TileCoordMNKL{}) == 4, "TileCoordMNKL must be rank 4");
// Separate out problem shape for convenience
auto M = get<0>(problem_shape_mnkl);
auto N = get<1>(problem_shape_mnkl);
auto L = get<3>(problem_shape_mnkl);
auto [M, N, K, L] = problem_shape_mnkl;
// Slice to get the tile this CTA is responsible for
auto [m_coord, n_coord, k_coord, l_coord] = cta_coord_mnkl;
bool is_accumulator_needed = K > 0;
if (is_accumulator_needed) {
// Wait for mma warp to fill tmem buffer with accumulator results
acc_pipeline.consumer_wait(acc_pipe_consumer_state);
}
// Batches are managed by using appropriate pointers to C and D matrices
auto problem_shape_mnl = append<3>(make_shape(M,N),Int<1>{});
@@ -320,13 +324,17 @@ public:
else {
Tensor tAcc = accumulators(make_coord(_,_),_0{},_0{}); // (CTA_M,CTA_N)
Tensor tTR_tAcc = thread_t2r.partition_S(tAcc); // (T2R,T2R_M,T2R_N)
copy(tiled_t2r, tTR_tAcc, tTR_rAcc);
if (is_accumulator_needed) {
copy(tiled_t2r, tTR_tAcc, tTR_rAcc);
} else {
fill(tTR_rAcc, 0);
}
}
if (is_accumulator_needed) {
cutlass::arch::fence_view_async_tmem_load();
acc_pipeline.consumer_release(acc_pipe_consumer_state);
++acc_pipe_consumer_state;
}
cutlass::arch::fence_view_async_tmem_load();
acc_pipeline.consumer_release(acc_pipe_consumer_state);
++acc_pipe_consumer_state;
// 2. Apply element-wise operation and store to gmem
// source is needed
@@ -628,12 +636,16 @@ public:
static_assert(rank(CtaCoordMNKL{}) == 4, "TileCoordMNKL must be rank 4");
static_assert(cute::sizeof_bits_v<ElementD> != 6, "Output element requires smem");
auto M = get<0>(problem_shape_mnkl);
auto N = get<1>(problem_shape_mnkl);
auto L = get<3>(problem_shape_mnkl);
// Separate out problem shape for convenience
auto [M, N, K, L] = problem_shape_mnkl;
// Slice to get the tile this CTA is responsible for
auto [m_coord, n_coord, k_coord, l_coord] = cta_coord_mnkl;
bool is_accumulator_needed = K > 0;
if (is_accumulator_needed) {
// Wait for mma warp to fill tmem buffer with accumulator results
acc_pipeline.consumer_wait(acc_pipe_consumer_state);
}
// Batches are managed by using appropriate pointers to C and D matrices
auto problem_shape_mnl = append<3>(make_shape(M,N),Int<1>{});
@@ -740,7 +752,7 @@ public:
auto synchronize = [] () CUTLASS_LAMBDA_FUNC_INLINE { cutlass::arch::NamedBarrier::sync(ThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); };
// The Epilogue Loop
auto epi_loop_fn = [&] (auto& cst_callbacks) CUTLASS_LAMBDA_FUNC_INLINE {
auto epi_loop_fn = [&] (auto& cst_callbacks, bool is_accumulator_needed) CUTLASS_LAMBDA_FUNC_INLINE {
// Ensure there are no threads from the previous wave writing to shared memory being utilized for the current wave.
synchronize();
cst_callbacks.begin();
@@ -802,10 +814,15 @@ public:
Tensor tTR_rAcc_frg = recast<Array<ElementAccumulator, FragmentSize>>(coalesce(tTR_rAcc));
copy(tiled_t2r, tTR_tAcc_mn, tTR_rAcc);
if (is_accumulator_needed) {
copy(tiled_t2r, tTR_tAcc_mn, tTR_rAcc);
}
else {
fill(tTR_rAcc, 0);
}
// After the last tmem load, signal that tmem buffer is consumed and empty
if (do_acc_release) {
if (do_acc_release && is_accumulator_needed) {
cutlass::arch::fence_view_async_tmem_load();
acc_pipeline.consumer_release(acc_pipe_consumer_state);
++acc_pipe_consumer_state;
@@ -886,7 +903,7 @@ public:
// BEGIN EPILOGUE
//
auto cst_callbacks = fusion_callbacks.template get_consumer_store_callbacks<RefSrc>(cst_args);
epi_loop_fn(cst_callbacks);
epi_loop_fn(cst_callbacks, is_accumulator_needed);
return cute::make_tuple(acc_pipe_consumer_state);
}
include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp
+10 -5
View File
@@ -666,6 +666,7 @@ public:
int thread_idx = threadIdx.x % ThreadCount;
int warp_idx = thread_idx / NumThreadsPerWarp;
[[maybe_unused]] int lane_idx = thread_idx % NumThreadsPerWarp;
bool is_accumulator_needed = K > 0;
// Check to see if tensormaps have been replaced in gmem
// Only the first epilogue warp needs to perform TMA related operations
@@ -824,7 +825,7 @@ public:
static_assert(not (DelayTmaStore and ReuseSmemC and StagesC <= StagesD), "This TMA epilogue configuration will deadlock");
// The Epilogue Loop
auto epi_loop_fn = [&] (auto& cst_callbacks) CUTLASS_LAMBDA_FUNC_INLINE {
auto epi_loop_fn = [&] (auto& cst_callbacks, bool is_accumulator_needed) CUTLASS_LAMBDA_FUNC_INLINE {
bool is_producer_load_needed = fusion_callbacks.is_producer_load_needed();
bool is_C_load_needed = is_source_supported && fusion_callbacks.is_C_load_needed();
@@ -929,7 +930,7 @@ public:
++load_wait_state;
}
if (is_first_iteration) {
if (is_first_iteration && is_accumulator_needed) {
// Wait for mma warp to fill tmem buffer with accumulator results
acc_pipeline.consumer_wait(acc_pipe_consumer_state, acc_wait_token);
}
@@ -946,11 +947,15 @@ public:
// Copy accumulator tile from tmem to register
if (issue_tmem_load) {
copy(tiled_t2r, tTR_tAcc_mn, tTR_rAcc);
if (is_accumulator_needed) {
copy(tiled_t2r, tTR_tAcc_mn, tTR_rAcc);
} else {
fill(tTR_rAcc, 0);
}
}
// After the last tmem load, signal that tmem buffer is consumed and empty
if (do_acc_release) {
if (do_acc_release && is_accumulator_needed) {
cutlass::arch::fence_view_async_tmem_load();
acc_pipeline.consumer_release(acc_pipe_consumer_state);
++acc_pipe_consumer_state;
@@ -1021,7 +1026,7 @@ public:
// BEGIN EPILOGUE
//
auto cst_callbacks = fusion_callbacks.template get_consumer_store_callbacks<RefSrc>(cst_args);
epi_loop_fn(cst_callbacks);
epi_loop_fn(cst_callbacks, is_accumulator_needed);
return cute::make_tuple(load_pipe_consumer_state, store_pipe_producer_state, acc_pipe_consumer_state);
}
include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp
+4 -2
View File
@@ -1082,8 +1082,10 @@ struct CollectiveMma<
}
}
else {
// Wait for tmem accumulator buffer to become empty with a flipped phase
accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);
if (k_tile_count > 0) {
// Wait for tmem accumulator buffer to become empty with a flipped phase
accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);
}
}
CUTLASS_PRAGMA_NO_UNROLL
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp
+11 -3
View File
@@ -1238,7 +1238,7 @@ public:
}
}();
auto accumulator = collective_mainloop.slice_accumulator(tmem_storage, acc_stage);
if (is_mma_leader_cta) {
if (is_mma_leader_cta && k_tile_count > 0) {
mainloop_pipe_consumer_state = collective_mainloop.mma(
cute::make_tuple(mainloop_pipeline, accumulator_pipeline),
cute::make_tuple(mainloop_pipe_consumer_state, accumulator_pipe_producer_state),
@@ -1249,7 +1249,9 @@ public:
);
accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);
}
++accumulator_pipe_producer_state;
if (k_tile_count > 0) {
++accumulator_pipe_producer_state;
}
// Fetch next work tile
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
@@ -1373,7 +1375,13 @@ public:
do_tail_load = true;
}
current_wave++;
// Relevant only for OverlappingAccum cases.
// Only increment the wave if the problem shape K dimension is not 0, otherwise accumulator will be skipped.
if constexpr (IsOverlappingAccum) {
if (size<2>(problem_shape_MNKL) > 0) {
current_wave++;
}
}
// Fetch the next work tile
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(