From 020c700e9782723ed6f61d8c4f13a364658dcb43 Mon Sep 17 00:00:00 2001 From: ANIKET SHIVAM <3268307+ANIKET-SHIVAM@users.noreply.github.com> Date: Tue, 4 Nov 2025 08:25:39 -0800 Subject: [PATCH] support for K=0 for sm100 GG (#2746) --- .../75_blackwell_grouped_gemm.cu | 12 ++--- .../75_blackwell_grouped_gemm_block_scaled.cu | 12 ++--- .../75_blackwell_grouped_gemm/CMakeLists.txt | 2 +- .../cutlass/epilogue/collective/detail.hpp | 3 -- .../sm100_epilogue_array_nosmem.hpp | 51 ++++++++++++------- ...100_epilogue_array_tma_warpspecialized.hpp | 15 ++++-- ..._blockscaled_mma_array_warpspecialized.hpp | 6 ++- .../sm100_gemm_array_tma_warpspecialized.hpp | 14 +++-- 8 files changed, 72 insertions(+), 43 deletions(-) diff --git a/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm.cu b/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm.cu index 84c42b91..e5a70d54 100644 --- a/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm.cu +++ b/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm.cu @@ -317,14 +317,14 @@ struct Options { int m = cmd_line_m; int n = cmd_line_n; int k = cmd_line_k; - if (m < 1) { - m = alignment * ((rand() % 64) + 1); + if (m < 0) { + m = alignment * ((rand() % 64)); } - if (n < 1) { - n = alignment * ((rand() % 64) + 1); + if (n < 0) { + n = alignment * ((rand() % 64)); } - if (k < 1) { - k = alignment * ((rand() % 64) + 1); + if (k < 0) { + k = alignment * ((rand() % 64)); } problem_sizes_host.push_back({m, n, k}); } diff --git a/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm_block_scaled.cu b/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm_block_scaled.cu index a18828e2..81bb737a 100644 --- a/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm_block_scaled.cu +++ b/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm_block_scaled.cu @@ -379,14 +379,14 @@ struct Options { int m = cmd_line_m; int n = cmd_line_n; int k = cmd_line_k; - if (m < 1) { - m = alignment * ((rand() % 64) + 1); + if (m < 0) { + m = alignment * ((rand() % 64)); } - if (n < 1) { - n = alignment * ((rand() % 64) + 1); + if (n < 0) { + n = alignment * ((rand() % 64)); } - if (k < 1) { - k = alignment * ((rand() % 64) + 1); + if (k < 0) { + k = alignment * ((rand() % 64)); } problem_sizes_host.push_back({m, n, k}); } diff --git a/examples/75_blackwell_grouped_gemm/CMakeLists.txt b/examples/75_blackwell_grouped_gemm/CMakeLists.txt index 8669eb35..6e0529a3 100644 --- a/examples/75_blackwell_grouped_gemm/CMakeLists.txt +++ b/examples/75_blackwell_grouped_gemm/CMakeLists.txt @@ -38,7 +38,7 @@ set(TEST_EPILOGUE --alpha=0.5 --beta=0.5 --iterations=0) set(TEST_EPILOGUE_LARGE_GROUP --alpha=1.5 --beta=2.0 --groups=50 --iterations=0) # Random problem sizes set(TEST_EPILOGUE_OP --beta=0.5 --iterations=1) # Random problem sizes -set(TEST_EPILOGUE_OP_LARGE_GROUP --alpha=1.5 --iterations=1) # Random problem sizes +set(TEST_EPILOGUE_OP_LARGE_GROUP --alpha=1.5 --groups=50 --iterations=1) # Random problem sizes set(TEST_FIXED --m=2048 --n=5120 --k=8192 --iterations=0) # Fixed problem sizes set(TEST_FIXED_LARGE_GROUP --m=2048 --n=512 --k=512 --groups=51 --iterations=0) # Fixed problem sizes diff --git a/include/cutlass/epilogue/collective/detail.hpp b/include/cutlass/epilogue/collective/detail.hpp index f93eff7b..407ea27d 100644 --- a/include/cutlass/epilogue/collective/detail.hpp +++ b/include/cutlass/epilogue/collective/detail.hpp @@ -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()( acc_pipeline, acc_pipe_consumer_state, diff --git a/include/cutlass/epilogue/collective/sm100_epilogue_array_nosmem.hpp b/include/cutlass/epilogue/collective/sm100_epilogue_array_nosmem.hpp index afdc4528..e2c99f19 100644 --- a/include/cutlass/epilogue/collective/sm100_epilogue_array_nosmem.hpp +++ b/include/cutlass/epilogue/collective/sm100_epilogue_array_nosmem.hpp @@ -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 != 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>(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(cst_args); - epi_loop_fn(cst_callbacks); + epi_loop_fn(cst_callbacks, is_accumulator_needed); return cute::make_tuple(acc_pipe_consumer_state); } diff --git a/include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp b/include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp index e9f06f24..852deb7f 100644 --- a/include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp +++ b/include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp @@ -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(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); } diff --git a/include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp b/include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp index cf9a1349..9ec563c7 100644 --- a/include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp +++ b/include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp @@ -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 diff --git a/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp b/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp index 185b9d5d..c03d67c0 100644 --- a/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp +++ b/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp @@ -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(