The execution order between GEMM0 of the next iteration and GEMM1 of the current iteration in Pingpong scheduling pipeline for overlapping gemms and softmax between warpgroups

[tengdecheng]

Hi, I have a question about the pipeline of Pingpong scheduling. I studied Figure 1 (Pingpong scheduling for 2 warpgroups to overlap softmax and GEMMs) and found that after warpgroup1 and warpgroup2 have completed softmax, warpgroup2 provides an arrive barrier signal to wargroup1. At this time, warpgroup1 will continue to execute GEMM1 and GEMM0.
Here I have a question about the order of the two: In Figure 1, GEMM1 (PV of the current iteration) takes precedence over the execution of QKT (GEMM0) of the next iteration.

But in the code, it seems to be the opposite that GEMM0 of the next iteration executed firstly and then GEMM1 of the current iteration executed, and it seems to be the following execution order:

warp_scheduler_barrier_sync();
flash::gemm</*zero_init=*/true, /*wg_wait=*/-1>(tiled_mma0, tSrQ, tSrK(_, _, _, smem_pipe_read_k.index()), tSrS);
softmax.rescale_o(tOrO, scores_scale);
consumer_wait(pipeline_v, smem_pipe_read_v);
flash::gemm</*zero_init=*/false, /*wg_wait=*/-1>(tiled_mma1, tOrP, tOrV(_, _, _, smem_pipe_read_v.index()), tOrO);
warp_scheduler_barrier_arrive();

I don’t find the implementation to guarantee the execution order of the two that GEMM1 of the current iteration is executed before GEMM0 of the next iteration. Could you please tell me where is it? Thank you very much.