Skip to content

fix: missing PDL wait on main_sf_load in sm103 blockscaled GEMM #3279

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tianyuxbear wants to merge 1 commit into
base: main
Choose a base branch
from
+3 −8
Open

fix: missing PDL wait on main_sf_load in sm103 blockscaled GEMM #3279

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 3 additions & 8 deletions include/cutlass/gemm/kernel/sm103_blockscaled_gemm_tma_warpspecialized.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -668,14 +668,14 @@ class GemmUniversal<
Tensor accumulators = cutlass::detail::make_sm100_accumulator<AccumulatorPipelineStageCount, IsOverlappingAccum>(
tiled_mma, acc_shape, EpilogueTile{});

// Ensure memory ops in this kernel are not done prior to completion of dependent grids.
cutlass::arch::wait_on_dependent_grids();

#if 1
pipeline_init_wait(cluster_size);

if (is_participant.main_ab_load) {
set_warpgroup_reg_dealloc();
// Ensure that the prefetched kernel does not touch
// unflushed global memory prior to this instruction
cutlass::arch::wait_on_dependent_grids();

bool do_load_order_arrive = is_epi_load_needed;
auto load_inputs = collective_mainloop.load_ab_init(
Expand Down Expand Up @@ -752,8 +752,6 @@ class GemmUniversal<
// why this variable is needed.
bool requires_clc_query = true;

cutlass::arch::wait_on_dependent_grids();

do {
if (requires_clc_query) {
// Throttle CLC query to mitigate workload imbalance caused by skews among persistent workers.
Expand Down Expand Up @@ -957,9 +955,6 @@ class GemmUniversal<
}
else if (not IsNoSmemEpilogue and is_participant.epi_load) {
set_warpgroup_reg_dealloc();
// Ensure that the prefetched kernel does not touch
// unflushed global memory prior to this instruction
cutlass::arch::wait_on_dependent_grids();

bool do_load_order_wait = true;
bool do_tail_load = false;
Expand Down