Skip to content
Merged
6 changes: 4 additions & 2 deletions aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py
Original file line number Diff line number Diff line change
Expand Up @@ -446,6 +446,7 @@ def kernel_gemm(
b0_frag = b_s2r.load(b_cur0, preshuffled=b_preshuffled)
a0_frag = a_s2r.load(a_cur0)
a_g2s.load(a_next1, A1_gl_offset + (k + 1) * BLOCK_K)
rocdl.s_waitcnt(lgkmcnt=0)
rocdl.s_barrier()

c00_frag = mfma.call(a0_frag, b0_frag, c00_frag)
Expand Down Expand Up @@ -477,7 +478,7 @@ def kernel_gemm(
k = K_ITERS - 2
b0_frag = b_s2r.load(b_cur0, preshuffled=b_preshuffled)
a0_frag = a_s2r.load(a_cur0)
rocdl.s_barrier()
wait_barrier(0)

c00_frag = mfma.call(a0_frag, b0_frag, c00_frag)

Expand Down Expand Up @@ -523,7 +524,8 @@ def kernel_gemm(
c10_frag = mfma.call(a1_frag, b0_frag, c10_frag, set_prio=False)
c11_frag = mfma.call(a1_frag, b1_frag, c11_frag, set_prio=False)
rocdl.s_setprio(0)
rocdl.s_barrier()
if wave_m == 0:
rocdl.s_barrier()

# Scale and store back to gmem
wave_n_offset = wave_n * (N_TILES_B * 16)
Expand Down